_aiter_ops.py · tduncan/Nvidia-Optimization-Stack-VLLM

@tduncan / Nvidia-Optimization-Stack-VLLM

Code Commits Issues Pulls Cloud Agent CI/CD

main

py 4,601 lines 149 KB

Raw

1	# SPDX-License-Identifier: Apache-2.0
2	# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3	import ctypes
4	import functools
5	import os
6	from collections.abc import Callable
7	from typing import TYPE_CHECKING, Any
8
9	import torch
10	import torch.distributed as dist
11	from torch._ops import OpOverload
12
13	import vllm.envs as envs
14	from vllm.logger import init_logger
15	from vllm.platforms import current_platform
16	from vllm.utils.torch_utils import direct_register_custom_op
17	from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
18	rocm_aiter_sparse_attn_indexer,
19	rocm_aiter_sparse_attn_indexer_fake,
20	)
21
22	if TYPE_CHECKING:
23	from aiter import ActivationType, QuantType
24
25	logger = init_logger(__name__)
26
27	# fp8_dtype is not cached.
28	# on ROCm the fp8_dtype always calls is_fp8_fnuz
29	# which is a host op, so we cache it once here.
30	FP8_DTYPE = current_platform.fp8_dtype()
31	_HIPB_MM_INITIALIZED_DEVICES: set[int] = set()
32	_FP8_E4M3_DTYPES = {
33	dtype
34	for dtype in (
35	getattr(torch, "float8_e4m3fn", None),
36	getattr(torch, "float8_e4m3fnuz", None),
37	)
38	if dtype is not None
39	}
40	KB = 1024
41	MB = 1024 * KB
42
43
44	def _is_fp8_e4m3_tensor(t: torch.Tensor \| None) -> bool:
45	return t is not None and t.dtype in _FP8_E4M3_DTYPES
46
47
48	def _get_or_create_aiter_qr_rmsnorm_comm(
49	device_comm: Any,
50	) -> tuple[int, int, bool] \| None:
51	"""Create the AITER QR communicator used by fused QR+RMSNorm.
52
53	vLLM's normal QuickReduce communicator owns vLLM C++ QR buffers. The fused
54	QR+RMSNorm kernel added in AITER needs an AITER DeviceComms instance, so we
55	lazily create one on the same TP CPU group and keep it on the vLLM device
56	communicator for reuse.
57	"""
58	existing = getattr(device_comm, "_aiter_qr_rmsnorm_comm", None)
59	if existing is False:
60	return None
61	if existing is not None:
62	return existing
63
64	try:
65	import aiter
66	except Exception:
67	device_comm._aiter_qr_rmsnorm_comm = False
68	return None
69
70	if not hasattr(aiter, "qr_all_reduce_rmsnorm"):
71	device_comm._aiter_qr_rmsnorm_comm = False
72	return None
73
74	group = getattr(device_comm, "cpu_group", None)
75	if group is None:
76	device_comm._aiter_qr_rmsnorm_comm = False
77	return None
78
79	world_size = dist.get_world_size(group=group)
80	rank = dist.get_rank(group=group)
81	if world_size not in (2, 4, 8):
82	device_comm._aiter_qr_rmsnorm_comm = False
83	return None
84
85	qr_max_size_mb = envs.VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB
86	qr_max_size = None if qr_max_size_mb is None else qr_max_size_mb * MB
87
88	try:
89	ptr = aiter.init_custom_qr(rank, world_size, qr_max_size)
90	handle = aiter.qr_get_handle(ptr)
91	handles = [None] * world_size
92	dist.all_gather_object(handles, handle, group=group)
93	aiter.qr_open_handles(ptr, handles)
94	except Exception:
95	device_comm._aiter_qr_rmsnorm_comm = False
96	return None
97
98	comm = (ptr, world_size, envs.VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16)
99	device_comm._aiter_qr_rmsnorm_comm = comm
100	return comm
101
102
103	def _ensure_hipb_mm_extension_initialized() -> None:
104	import aiter
105
106	device = torch.accelerator.current_device_index()
107	if device not in _HIPB_MM_INITIALIZED_DEVICES:
108	aiter.hipb_create_extension()
109	_HIPB_MM_INITIALIZED_DEVICES.add(device)
110
111
112	def is_aiter_found() -> bool:
113	from importlib.util import find_spec
114
115	return find_spec("aiter") is not None
116
117
118	# `find_spec` is not torch.compile compatible.
119	# In cases where aiter availability might have
120	# been checked in forward passes that are torch compiled.
121	# we keep this global outside to not cause torch compile breaks.
122	IS_AITER_FOUND = is_aiter_found()
123
124
125	class _DlInfo(ctypes.Structure):
126	_fields_ = [
127	("dli_fname", ctypes.c_char_p),
128	("dli_fbase", ctypes.c_void_p),
129	("dli_sname", ctypes.c_char_p),
130	("dli_saddr", ctypes.c_void_p),
131	]
132
133
134	_HIP_RUNTIME_HANDLE: ctypes.CDLL \| None = None
135
136
137	def _find_torch_hip_runtime_path() -> str \| None:
138	"""Find the HIP runtime already selected by PyTorch's extension graph."""
139	try:
140	torch_extension = ctypes.CDLL(torch._C.__file__)
141	hip_symbol = torch_extension.hipGetDeviceCount
142	dladdr = ctypes.CDLL(None).dladdr
143	dladdr.argtypes = [ctypes.c_void_p, ctypes.POINTER(_DlInfo)]
144	dladdr.restype = ctypes.c_int
145	info = _DlInfo()
146	if not dladdr(ctypes.cast(hip_symbol, ctypes.c_void_p), ctypes.byref(info)):
147	return None
148	except (AttributeError, OSError, TypeError, ctypes.ArgumentError):
149	return None
150	if info.dli_fname is None:
151	return None
152
153	runtime_path = os.fsdecode(info.dli_fname)
154	if "libamdhip64.so" not in os.path.basename(runtime_path):
155	return None
156	return runtime_path
157
158
159	def _promote_torch_hip_runtime() -> None:
160	"""Make PyTorch's HIP runtime symbols visible to AITER JIT extensions.
161
162	PyTorch loads its native extension with ``RTLD_LOCAL``. Some AITER JIT
163	extensions rely on HIP registration and launch symbols being available in
164	the process-global lookup scope. Reopening the exact runtime that provides
165	PyTorch's ``hipGetDeviceCount`` symbol promotes the existing DSO instead of
166	selecting or initializing a second HIP runtime.
167	"""
168	global _HIP_RUNTIME_HANDLE
169	if _HIP_RUNTIME_HANDLE is not None:
170	return
171
172	runtime_path = _find_torch_hip_runtime_path()
173	if runtime_path is None:
174	logger.warning_once(
175	"Could not locate PyTorch's HIP runtime; AITER JIT extensions may "
176	"fail to resolve HIP symbols."
177	)
178	return
179
180	try:
181	# Retain the handle for the process lifetime. Besides preventing a
182	# future dlclose, this documents that the RTLD_GLOBAL promotion is
183	# intentional and must outlive every lazily imported AITER extension.
184	_HIP_RUNTIME_HANDLE = ctypes.CDLL(runtime_path, mode=ctypes.RTLD_GLOBAL)
185	except OSError as err:
186	logger.warning_once(
187	"Could not promote PyTorch's HIP runtime at %s; AITER JIT "
188	"extensions may fail to resolve HIP symbols: %s",
189	runtime_path,
190	err,
191	)
192
193
194	def _maybe_promote_torch_hip_runtime() -> None:
195	if current_platform.is_rocm() and IS_AITER_FOUND:
196	_promote_torch_hip_runtime()
197
198
199	_maybe_promote_torch_hip_runtime()
200
201
202	def is_aiter_found_and_supported() -> bool:
203	"""Check if AITER library is available and platform supports it.
204
205	Checks: platform (ROCm), device arch is CDNA 3 or better, and library existence.
206	Does NOT check environment variables - that's handled by rocm_aiter_ops.is_enabled().
207
208	This function determines if aiter CAN be used, not if it SHOULD be used.
209
210	Separation of concerns:
211	- This function: Can aiter work on this system? (platform + library availability)
212	- rocm_aiter_ops.is_enabled(): Should aiter be used by default? (adds env var check)
213	- Backend selection: Can explicitly request aiter regardless of env var
214
215	This allows explicit backend selection via attention_config to work even when
216	VLLM_ROCM_USE_AITER=0, while preventing unwanted JIT warnings for auto-discovery.
217	"""
218	if current_platform.is_rocm() and IS_AITER_FOUND:
219	from vllm.platforms.rocm import get_cdna_version
220
221	return get_cdna_version() > 2
222	return False
223
224
225	def is_aiter_found_and_supported_on_rdna4() -> bool:
226	"""RDNA4 (gfx12) analog of `is_aiter_found_and_supported()`.
227
228	gfx12 has no aiter CK build, so this deliberately stays off the gfx9
229	`@if_aiter_supported` umbrella; it reports only that aiter's Triton
230	kernels are usable here. Like its gfx9 counterpart it checks platform +
231	arch + library availability and does not check environment variables.
232	"""
233	if current_platform.is_rocm() and IS_AITER_FOUND:
234	from vllm.platforms.rocm import on_rdna4
235
236	return on_rdna4()
237	return False
238
239
240	def _triton_gemm_config_is_tuned(config_name: str, N: int, K: int) -> bool:
241	from aiter.ops.triton.utils.gemm_config_utils import get_gemm_config
242
243	# any M shape under 1024 should be tuned
244	M_TUNE_PROBE = 128
245
246	return get_gemm_config(config_name, M_TUNE_PROBE, N, K)[1]
247
248
249	def _ck_gemm_shape_is_tuned(
250	N: int, K: int, q_dtype_w: torch.dtype, csv_attr: str
251	) -> bool:
252	l_m = (
253	[1, 2, 4]
254	+ list(range(8, 513, 8))
255	+ [1024, 1536]
256	+ [2**i for i in range(11, 19)]
257	)
258	try:
259	from aiter.ops.gemm_op_a8w8 import (
260	AITER_CONFIGS,
261	get_GEMM_config_with_quant_type,
262	)
263
264	csv_path = getattr(AITER_CONFIGS, csv_attr)
265	return any(
266	get_GEMM_config_with_quant_type(M, N, K, q_dtype_w, csv_path) is not None
267	for M in l_m
268	)
269	except (AttributeError, ImportError, OSError):
270	logger.warning_once(
271	"Could not read aiter CK GEMM configs from AITER_CONFIGS.%s; "
272	"treating all shapes as untuned.",
273	csv_attr,
274	)
275	return False
276
277
278	def if_aiter_supported(func: Callable) -> Callable:
279	"""Decorator that only executes the function if
280	ROCm AITER package is supported and enabled on gfx9 archs.
281	"""
282
283	@functools.wraps(func)
284	def wrapper(args, *kwargs):
285	if is_aiter_found_and_supported():
286	return func(args, *kwargs)
287
288	return None
289
290	return wrapper
291
292
293	def _validate_rocm_aiter_fused_moe_shared_expert_args(
294	shared_w1: torch.Tensor \| None,
295	shared_w2: torch.Tensor \| None,
296	shared_w1_scale: torch.Tensor \| None,
297	shared_w2_scale: torch.Tensor \| None,
298	shared_expert_id: int,
299	) -> bool:
300	shared_tensors = (shared_w1, shared_w2, shared_w1_scale, shared_w2_scale)
301	has_shared_expert = any(tensor is not None for tensor in shared_tensors)
302	if has_shared_expert:
303	if not all(tensor is not None for tensor in shared_tensors):
304	raise ValueError(
305	"Heterogeneous fused MoE requires both shared weights and scales."
306	)
307	if shared_expert_id < 0:
308	raise ValueError(
309	"Heterogeneous fused MoE requires a non-negative shared expert ID."
310	)
311	elif shared_expert_id >= 0:
312	raise ValueError(
313	"A non-negative shared expert ID requires shared weights and scales."
314	)
315	return has_shared_expert
316
317
318	def _rocm_aiter_fused_moe_impl(
319	hidden_states: torch.Tensor,
320	w1: torch.Tensor,
321	w2: torch.Tensor,
322	topk_weight: torch.Tensor,
323	topk_ids: torch.Tensor,
324	expert_mask: torch.Tensor \| None = None,
325	activation_method: int = 0,
326	quant_method: int = 0,
327	doweight_stage1: bool = False,
328	w1_scale: torch.Tensor \| None = None,
329	w2_scale: torch.Tensor \| None = None,
330	a1_scale: torch.Tensor \| None = None,
331	a2_scale: torch.Tensor \| None = None,
332	num_local_tokens: torch.Tensor \| None = None,
333	output_dtype: torch.dtype \| None = None,
334	hidden_pad: int = 0,
335	intermediate_pad: int = 0,
336	gate_mode: str = "",
337	bias1: torch.Tensor \| None = None,
338	bias2: torch.Tensor \| None = None,
339	moe_sorting_dispatch_policy: int = 0,
340	swiglu_limit: float = 0.0,
341	beta: float \| None = None,
342	linear_beta: float \| None = None,
343	shared_w1: torch.Tensor \| None = None,
344	shared_w2: torch.Tensor \| None = None,
345	shared_w1_scale: torch.Tensor \| None = None,
346	shared_w2_scale: torch.Tensor \| None = None,
347	shared_expert_id: int = -1,
348	q_dtype_a: torch.dtype \| None = None,
349	) -> torch.Tensor:
350	has_shared_expert = _validate_rocm_aiter_fused_moe_shared_expert_args(
351	shared_w1,
352	shared_w2,
353	shared_w1_scale,
354	shared_w2_scale,
355	shared_expert_id,
356	)
357
358	from aiter import ActivationType, QuantType
359	from aiter.fused_moe import fused_moe
360
361	activation = ActivationType(activation_method)
362	quant_type = QuantType(quant_method)
363
364	extra_kwargs: dict = {}
365	if gate_mode and rocm_aiter_ops.fused_moe_supports_gate_mode():
366	extra_kwargs["gate_mode"] = gate_mode
367	if (
368	getattr(ActivationType, "Situv2", None) is not None
369	and activation == ActivationType.Situv2
370	):
371	extra_kwargs["beta"] = beta
372	extra_kwargs["linear_beta"] = linear_beta
373
374	if has_shared_expert:
375	extra_kwargs.update(
376	shared_w1=shared_w1,
377	shared_w2=shared_w2,
378	shared_w1_scale=shared_w1_scale,
379	shared_w2_scale=shared_w2_scale,
380	shared_expert_id=shared_expert_id,
381	)
382	if q_dtype_a is not None:
383	# DeepSeek V4.1 a4w4 override (use_mxfp4_w4a4_dsv4 in
384	# rocm_aiter_moe.py). rocm_aiter_ops.fused_moe_supports_quant_dtype_a()
385	# is checked at config time, so this is only reached on an AITER
386	# build new enough to accept it.
387	extra_kwargs["quant_dtype_a"] = q_dtype_a
388
389	return fused_moe(
390	hidden_states,
391	w1,
392	w2,
393	topk_weight,
394	topk_ids,
395	expert_mask,
396	activation,
397	quant_type,
398	doweight_stage1,
399	w1_scale,
400	w2_scale,
401	a1_scale,
402	a2_scale,
403	num_local_tokens=num_local_tokens,
404	dtype=output_dtype,
405	hidden_pad=hidden_pad,
406	intermediate_pad=intermediate_pad,
407	bias1=bias1,
408	bias2=bias2,
409	moe_sorting_dispatch_policy=moe_sorting_dispatch_policy,
410	swiglu_limit=swiglu_limit,
411	**extra_kwargs,
412	)
413
414
415	def _rocm_aiter_fused_moe_fake(
416	hidden_states: torch.Tensor,
417	w1: torch.Tensor,
418	w2: torch.Tensor,
419	topk_weight: torch.Tensor,
420	topk_ids: torch.Tensor,
421	expert_mask: torch.Tensor \| None = None,
422	activation_method: int = 0,
423	quant_method: int = 0,
424	doweight_stage1: bool = False,
425	w1_scale: torch.Tensor \| None = None,
426	w2_scale: torch.Tensor \| None = None,
427	a1_scale: torch.Tensor \| None = None,
428	a2_scale: torch.Tensor \| None = None,
429	num_local_tokens: torch.Tensor \| None = None,
430	output_dtype: torch.dtype \| None = None,
431	hidden_pad: int = 0,
432	intermediate_pad: int = 0,
433	gate_mode: str = "",
434	bias1: torch.Tensor \| None = None,
435	bias2: torch.Tensor \| None = None,
436	moe_sorting_dispatch_policy: int = 0,
437	swiglu_limit: float = 0.0,
438	beta: float \| None = None,
439	linear_beta: float \| None = None,
440	shared_w1: torch.Tensor \| None = None,
441	shared_w2: torch.Tensor \| None = None,
442	shared_w1_scale: torch.Tensor \| None = None,
443	shared_w2_scale: torch.Tensor \| None = None,
444	shared_expert_id: int = -1,
445	q_dtype_a: torch.dtype \| None = None,
446	) -> torch.Tensor:
447	if output_dtype is not None:
448	return torch.empty_like(hidden_states, dtype=output_dtype)
449	return torch.empty_like(hidden_states)
450
451
452	def _rocm_aiter_asm_moe_tkw1_impl(
453	hidden_states: torch.Tensor,
454	w1: torch.Tensor,
455	w2: torch.Tensor,
456	topk_weights: torch.Tensor,
457	topk_ids: torch.Tensor,
458	fc1_scale: torch.Tensor \| None = None,
459	fc2_scale: torch.Tensor \| None = None,
460	fc1_smooth_scale: torch.Tensor \| None = None,
461	fc2_smooth_scale: torch.Tensor \| None = None,
462	a16: bool = False,
463	per_tensor_quant_scale: torch.Tensor \| None = None,
464	expert_mask: torch.Tensor \| None = None,
465	activation_method: int = 0,
466	) -> torch.Tensor:
467	from aiter import ActivationType
468	from aiter.fused_moe_bf16_asm import asm_moe_tkw1
469
470	activation = ActivationType(activation_method)
471
472	return asm_moe_tkw1(
473	hidden_states,
474	w1,
475	w2,
476	topk_weights,
477	topk_ids,
478	fc1_scale=fc1_scale,
479	fc2_scale=fc2_scale,
480	fc1_smooth_scale=fc1_smooth_scale,
481	fc2_smooth_scale=fc2_smooth_scale,
482	a16=a16,
483	per_tensor_quant_scale=per_tensor_quant_scale,
484	expert_mask=expert_mask,
485	activation=activation,
486	)
487
488
489	def _rocm_aiter_asm_moe_tkw1_fake(
490	hidden_states: torch.Tensor,
491	w1: torch.Tensor,
492	w2: torch.Tensor,
493	topk_weights: torch.Tensor,
494	topk_ids: torch.Tensor,
495	fc1_scale: torch.Tensor \| None = None,
496	fc2_scale: torch.Tensor \| None = None,
497	fc1_smooth_scale: torch.Tensor \| None = None,
498	fc2_smooth_scale: torch.Tensor \| None = None,
499	a16: bool = False,
500	per_tensor_quant_scale: torch.Tensor \| None = None,
501	expert_mask: torch.Tensor \| None = None,
502	activation_method: int = 0,
503	) -> torch.Tensor:
504	return torch.empty_like(hidden_states)
505
506
507	def _rocm_aiter_topk_softmax_impl(
508	topk_weights: torch.Tensor,
509	topk_indices: torch.Tensor,
510	token_expert_indices: torch.Tensor,
511	gating_output: torch.Tensor,
512	renormalize: bool,
513	num_shared_experts: int = 0,
514	shared_expert_scoring_func: str = "",
515	) -> None:
516	from aiter import topk_softmax
517
518	topk_softmax(
519	topk_weights,
520	topk_indices,
521	token_expert_indices,
522	gating_output,
523	renormalize,
524	num_shared_experts,
525	shared_expert_scoring_func,
526	)
527
528
529	def _rocm_aiter_topk_sigmoid_impl(
530	topk_weights: torch.Tensor,
531	topk_indices: torch.Tensor,
532	gating_output: torch.Tensor,
533	) -> None:
534	from aiter import topk_sigmoid
535
536	topk_sigmoid(topk_weights, topk_indices, gating_output)
537
538
539	def _rocm_aiter_biased_grouped_topk_impl(
540	gating_output: torch.Tensor,
541	correction_bias: torch.Tensor,
542	topk_weights: torch.Tensor,
543	topk_ids: torch.Tensor,
544	num_expert_group: int,
545	topk_group: int,
546	need_renorm: bool,
547	routed_scaling_factor: float = 1.0, # mul to topk_weights
548	) -> None:
549	from aiter import biased_grouped_topk
550
551	biased_grouped_topk(
552	gating_output,
553	correction_bias,
554	topk_weights,
555	topk_ids,
556	num_expert_group,
557	topk_group,
558	need_renorm,
559	routed_scaling_factor,
560	)
561
562
563	def _rocm_aiter_grouped_topk_impl(
564	gating_output: torch.Tensor,
565	topk_weights: torch.Tensor,
566	topk_ids: torch.Tensor,
567	num_expert_group: int,
568	topk_group: int,
569	need_renorm: bool,
570	scoring_func: str = "softmax",
571	routed_scaling_factor: float = 1.0, # mul to topk_weights
572	) -> None:
573	is_softmax = scoring_func == "softmax"
574	from aiter import grouped_topk
575
576	grouped_topk(
577	gating_output,
578	topk_weights,
579	topk_ids,
580	num_expert_group,
581	topk_group,
582	need_renorm,
583	is_softmax,
584	routed_scaling_factor,
585	)
586
587
588	def _rocm_aiter_fused_topk_impl(
589	x: torch.Tensor,
590	router_logits: torch.Tensor,
591	top_k: int,
592	gate_up: bool,
593	) -> tuple[torch.Tensor, torch.Tensor]:
594	from aiter.fused_moe import fused_topk
595
596	# fused_topk returns (topk_weights, topk_indices)
597	return fused_topk(x, router_logits, top_k, gate_up)
598
599
600	def _rocm_aiter_fused_topk_fake(
601	x: torch.Tensor,
602	router_logits: torch.Tensor,
603	top_k: int,
604	gate_up: bool,
605	) -> tuple[torch.Tensor, torch.Tensor]:
606	num_tokens = x.shape[0]
607	topk_weights = torch.empty(
608	(num_tokens, top_k), dtype=torch.float32, device=x.device
609	)
610	topk_indices = torch.empty((num_tokens, top_k), dtype=torch.int32, device=x.device)
611	return topk_weights, topk_indices
612
613
614	# Cache whether aiter supports FP8 MLA parameters
615	_AITER_MLA_SUPPORTS_FP8: bool \| None = None
616
617
618	def _check_aiter_mla_fp8_support() -> bool:
619	"""Check if aiter.mla.mla_decode_fwd supports q_scale and kv_scale parameters."""
620	global _AITER_MLA_SUPPORTS_FP8
621	if _AITER_MLA_SUPPORTS_FP8 is None:
622	try:
623	import inspect
624
625	from aiter.mla import mla_decode_fwd
626
627	sig = inspect.signature(mla_decode_fwd)
628	_AITER_MLA_SUPPORTS_FP8 = (
629	"q_scale" in sig.parameters and "kv_scale" in sig.parameters
630	)
631	except (
632	ImportError,
633	ModuleNotFoundError,
634	AttributeError,
635	ValueError,
636	TypeError,
637	):
638	# ImportError/ModuleNotFoundError: aiter.mla module not available
639	# AttributeError: mla_decode_fwd doesn't exist
640	# ValueError: mla_decode_fwd has no signature (e.g., built-in)
641	# TypeError: mla_decode_fwd is not a callable
642	_AITER_MLA_SUPPORTS_FP8 = False
643	return _AITER_MLA_SUPPORTS_FP8
644
645
646	def _rocm_aiter_mla_decode_fwd_impl(
647	q: torch.Tensor,
648	kv_buffer: torch.Tensor,
649	o: torch.Tensor,
650	qo_indptr: torch.Tensor,
651	max_seqlen_qo: int,
652	kv_indptr: torch.Tensor \| None = None,
653	kv_indices: torch.Tensor \| None = None,
654	kv_last_page_lens: torch.Tensor \| None = None,
655	sm_scale: float = 1.0,
656	logit_cap: float = 0.0,
657	q_scale: torch.Tensor \| None = None,
658	kv_scale: torch.Tensor \| None = None,
659	work_meta_data: torch.Tensor \| None = None,
660	work_indptr: torch.Tensor \| None = None,
661	work_info_set: torch.Tensor \| None = None,
662	reduce_indptr: torch.Tensor \| None = None,
663	reduce_final_map: torch.Tensor \| None = None,
664	reduce_partial_map: torch.Tensor \| None = None,
665	causal: bool = True,
666	) -> None:
667	from aiter.mla import mla_decode_fwd
668
669	kwargs: dict[str, float \| torch.Tensor \| None \| bool] = {
670	"sm_scale": sm_scale,
671	"logit_cap": logit_cap,
672	}
673
674	# Only pass q_scale and kv_scale if the aiter library supports them
675	if _check_aiter_mla_fp8_support():
676	kwargs["q_scale"] = q_scale
677	kwargs["kv_scale"] = kv_scale
678
679	kwargs["causal"] = causal
680
681	if work_meta_data is not None:
682	assert work_indptr is not None, (
683	"work_indptr must be provided with work_meta_data"
684	)
685	assert work_info_set is not None, (
686	"work_info_set must be provided with work_meta_data"
687	)
688	assert reduce_indptr is not None, (
689	"reduce_indptr must be provided with work_meta_data"
690	)
691	assert reduce_final_map is not None, (
692	"reduce_final_map must be provided with work_meta_data"
693	)
694	assert reduce_partial_map is not None, (
695	"reduce_partial_map must be provided with work_meta_data"
696	)
697	kwargs["work_meta_data"] = work_meta_data
698	kwargs["work_indptr"] = work_indptr
699	kwargs["work_info_set"] = work_info_set
700	kwargs["reduce_indptr"] = reduce_indptr
701	kwargs["reduce_final_map"] = reduce_final_map
702	kwargs["reduce_partial_map"] = reduce_partial_map
703
704	mla_decode_fwd(
705	q,
706	kv_buffer.view(-1, 1, 1, q.shape[-1]),
707	o,
708	qo_indptr,
709	kv_indptr,
710	kv_indices,
711	kv_last_page_lens,
712	max_seqlen_qo,
713	**kwargs,
714	)
715
716
717	def _rocm_aiter_mla_decode_fwd_lse_impl(
718	q: torch.Tensor,
719	kv_buffer: torch.Tensor,
720	o: torch.Tensor,
721	qo_indptr: torch.Tensor,
722	max_seqlen_qo: int,
723	kv_indptr: torch.Tensor \| None = None,
724	kv_indices: torch.Tensor \| None = None,
725	kv_last_page_lens: torch.Tensor \| None = None,
726	sm_scale: float = 1.0,
727	logit_cap: float = 0.0,
728	q_scale: torch.Tensor \| None = None,
729	kv_scale: torch.Tensor \| None = None,
730	) -> torch.Tensor:
731	"""Run non-persistent AITER MLA decode and return natural-log LSE.
732
733	gfx942 persistent MLA kernels do not provide an LSE code object. Keeping
734	this as a separate op makes that constraint structural: callers that need
735	LSE cannot accidentally forward persistent work metadata.
736	"""
737	from aiter.mla import get_meta_param, mla_decode_fwd
738
739	kwargs: dict[str, float \| int \| torch.Tensor \| None \| bool] = {
740	"sm_scale": sm_scale,
741	"logit_cap": logit_cap,
742	"return_lse": True,
743	}
744	if _check_aiter_mla_fp8_support():
745	kwargs["q_scale"] = q_scale
746	kwargs["kv_scale"] = kv_scale
747
748	if (
749	q.dtype == torch.bfloat16
750	and kv_buffer.dtype == torch.bfloat16
751	and q.shape[1] == 32
752	):
753	from vllm.platforms.rocm import on_gfx950
754
755	if on_gfx950():
756	# The gfx950 H32 BF16 kernel mishandles ragged tails with split KV.
757	# Its single-split path returns correct output and natural-log LSE.
758	kwargs["num_kv_splits"] = 1
759	kwargs["num_kv_splits_indptr"] = torch.arange(
760	qo_indptr.shape[0], dtype=torch.int32, device=q.device
761	)
762
763	if q.dtype == FP8_DTYPE:
764	assert kv_indices is not None
765	num_kv_splits, num_kv_splits_indptr = get_meta_param(
766	None,
767	qo_indptr.shape[0] - 1,
768	kv_indices.numel(),
769	q.shape[1],
770	max_seqlen_qo,
771	q.dtype,
772	)
773	if num_kv_splits == 1:
774	# gfx942's one-split FP8 asm writes the final output directly but
775	# does not write either LSE buffer. Force the normal split reducer,
776	# which produces both the same output and an accurate natural LSE.
777	num_kv_splits = 2
778	num_kv_splits_indptr = torch.arange(
779	0,
780	(qo_indptr.shape[0]) * num_kv_splits,
781	num_kv_splits,
782	dtype=torch.int32,
783	device=q.device,
784	)
785	kwargs["num_kv_splits"] = num_kv_splits
786	kwargs["num_kv_splits_indptr"] = num_kv_splits_indptr
787
788	_, final_lse = mla_decode_fwd(
789	q,
790	kv_buffer.view(-1, 1, 1, q.shape[-1]),
791	o,
792	qo_indptr,
793	kv_indptr,
794	kv_indices,
795	kv_last_page_lens,
796	max_seqlen_qo,
797	**kwargs,
798	)
799	if final_lse is None:
800	raise RuntimeError("AITER MLA decode did not return the requested LSE")
801	return final_lse
802
803
804	def _rocm_aiter_mla_decode_fwd_lse_fake(
805	q: torch.Tensor,
806	kv_buffer: torch.Tensor,
807	o: torch.Tensor,
808	qo_indptr: torch.Tensor,
809	max_seqlen_qo: int,
810	kv_indptr: torch.Tensor \| None = None,
811	kv_indices: torch.Tensor \| None = None,
812	kv_last_page_lens: torch.Tensor \| None = None,
813	sm_scale: float = 1.0,
814	logit_cap: float = 0.0,
815	q_scale: torch.Tensor \| None = None,
816	kv_scale: torch.Tensor \| None = None,
817	) -> torch.Tensor:
818	return torch.empty(q.shape[:-1], dtype=torch.float32, device=q.device)
819
820
821	def _rocm_aiter_w8a8_gemm_impl(
822	A: torch.Tensor,
823	B: torch.Tensor,
824	As: torch.Tensor,
825	Bs: torch.Tensor,
826	bias: torch.Tensor \| None = None,
827	output_dtype: torch.dtype = torch.float16,
828	) -> torch.Tensor:
829	from aiter import gemm_a8w8_CK
830
831	# gemm_a8w8_CK(a, b, scale_a, scale_b, bias) expects
832	# a to be [M, K]
833	# b to be [N, K]
834	# CutlassInt8ScaledMMLinearKernel prepare weight `w_q` in [K, N] format
835	return gemm_a8w8_CK(A, B, As, Bs, bias, output_dtype)
836
837
838	def _rocm_aiter_w8a8_gemm_fake(
839	A: torch.Tensor,
840	B: torch.Tensor,
841	As: torch.Tensor,
842	Bs: torch.Tensor,
843	bias: torch.Tensor \| None = None,
844	output_dtype: torch.dtype = torch.float16,
845	) -> torch.Tensor:
846	m = A.shape[0]
847	n = B.shape[0]
848	Y = torch.empty(m, n, dtype=output_dtype, device=A.device)
849	return Y
850
851
852	def _rocm_aiter_preshuffled_per_token_w8a8_gemm_impl(
853	A: torch.Tensor,
854	B: torch.Tensor,
855	As: torch.Tensor,
856	Bs: torch.Tensor,
857	bias: torch.Tensor \| None = None,
858	output_dtype: torch.dtype = torch.float16,
859	) -> torch.Tensor:
860	from aiter import gemm_a8w8_bpreshuffle
861
862	output = gemm_a8w8_bpreshuffle(A, B, As, Bs, None, output_dtype)
863	if bias is not None:
864	output.add_(bias)
865	return output
866
867
868	def _rocm_aiter_preshuffled_per_token_w8a8_gemm_fake(
869	A: torch.Tensor,
870	B: torch.Tensor,
871	As: torch.Tensor,
872	Bs: torch.Tensor,
873	bias: torch.Tensor \| None = None,
874	output_dtype: torch.dtype = torch.float16,
875	) -> torch.Tensor:
876	m = A.shape[0]
877	n = B.shape[0]
878	return torch.empty(m, n, dtype=output_dtype, device=A.device)
879
880
881	def _rocm_aiter_hipb_mm_fp8_impl(
882	A: torch.Tensor,
883	B: torch.Tensor,
884	As: torch.Tensor,
885	Bs: torch.Tensor,
886	bias: torch.Tensor \| None = None,
887	output_dtype: torch.dtype = torch.bfloat16,
888	) -> torch.Tensor:
889	from aiter import hipb_mm
890
891	_ensure_hipb_mm_extension_initialized()
892	return hipb_mm(
893	A,
894	B,
895	solution_index=-1,
896	bias=bias,
897	out_dtype=output_dtype,
898	scaleA=As,
899	scaleB=Bs,
900	scaleOut=None,
901	bpreshuffle=True,
902	)
903
904
905	def _rocm_aiter_hipb_mm_fp8_fake(
906	A: torch.Tensor,
907	B: torch.Tensor,
908	As: torch.Tensor,
909	Bs: torch.Tensor,
910	bias: torch.Tensor \| None = None,
911	output_dtype: torch.dtype = torch.bfloat16,
912	) -> torch.Tensor:
913	m = A.shape[0]
914	n = B.shape[1]
915	return torch.empty(m, n, dtype=output_dtype, device=A.device)
916
917
918	def _rocm_aiter_triton_gemm_a8w8_blockscale_impl(
919	A: torch.Tensor,
920	B: torch.Tensor,
921	As: torch.Tensor,
922	Bs: torch.Tensor,
923	output_dtype: torch.dtype = torch.float16,
924	) -> torch.Tensor:
925	from aiter.ops.triton.gemm_a8w8_blockscale import gemm_a8w8_blockscale
926
927	return gemm_a8w8_blockscale(A, B, As, Bs, dtype=output_dtype)
928
929
930	def _rocm_aiter_triton_gemm_a8w8_blockscale_fake(
931	A: torch.Tensor,
932	B: torch.Tensor,
933	As: torch.Tensor,
934	Bs: torch.Tensor,
935	output_dtype: torch.dtype = torch.float16,
936	) -> torch.Tensor:
937	m = A.shape[0]
938	n = B.shape[0]
939	Y = torch.empty(m, n, dtype=output_dtype, device=A.device)
940	return Y
941
942
943	def _rocm_aiter_gemm_a8w8_blockscale_impl(
944	A: torch.Tensor,
945	B: torch.Tensor,
946	As: torch.Tensor,
947	Bs: torch.Tensor,
948	output_dtype: torch.dtype = torch.float16,
949	) -> torch.Tensor:
950	from aiter import gemm_a8w8_blockscale
951
952	return gemm_a8w8_blockscale(A, B, As, Bs, dtype=output_dtype)
953
954
955	def _rocm_aiter_gemm_a8w8_blockscale_fake(
956	A: torch.Tensor,
957	B: torch.Tensor,
958	As: torch.Tensor,
959	Bs: torch.Tensor,
960	output_dtype: torch.dtype = torch.float16,
961	) -> torch.Tensor:
962	m = A.shape[0]
963	n = B.shape[0]
964	Y = torch.empty(m, n, dtype=output_dtype, device=A.device)
965	return Y
966
967
968	def _rocm_aiter_gemm_a8w8_blockscale_bpreshuffle_impl(
969	A: torch.Tensor,
970	B: torch.Tensor,
971	As: torch.Tensor,
972	Bs: torch.Tensor,
973	output_dtype: torch.dtype = torch.float16,
974	) -> torch.Tensor:
975	from aiter import gemm_a8w8_blockscale_bpreshuffle
976
977	# B preshuffled (shuffle_weight (16,16)); As column-major group scale.
978	return gemm_a8w8_blockscale_bpreshuffle(A, B, As, Bs, dtype=output_dtype)
979
980
981	def _rocm_aiter_gemm_a8w8_blockscale_bpreshuffle_fake(
982	A: torch.Tensor,
983	B: torch.Tensor,
984	As: torch.Tensor,
985	Bs: torch.Tensor,
986	output_dtype: torch.dtype = torch.float16,
987	) -> torch.Tensor:
988	m = A.shape[0]
989	n = B.shape[0]
990	Y = torch.empty(m, n, dtype=output_dtype, device=A.device)
991	return Y
992
993
994	def _rocm_aiter_rmsnorm_fused_add_dynamic_quant_impl(
995	x: torch.Tensor,
996	residual: torch.Tensor,
997	weight: torch.Tensor,
998	epsilon: float,
999	quant_dtype: torch.dtype,
1000	) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
1001	import aiter as rocm_aiter
1002
1003	assert quant_dtype in [torch.int8, FP8_DTYPE]
1004
1005	y_scale = torch.empty(x.shape[0], 1, dtype=torch.float32, device=x.device)
1006	out = torch.empty(x.shape, dtype=quant_dtype, device=x.device)
1007	residual_out = torch.empty_like(x)
1008
1009	rocm_aiter.rmsnorm2d_fwd_with_add_dynamicquant(
1010	out,
1011	x,
1012	residual,
1013	residual_out,
1014	y_scale,
1015	weight,
1016	epsilon,
1017	use_model_sensitive_rmsnorm=0,
1018	)
1019
1020	return out, residual_out, y_scale
1021
1022
1023	def _rocm_aiter_rmsnorm_fused_add_dynamic_quant_fake(
1024	x: torch.Tensor,
1025	residual: torch.Tensor,
1026	weight: torch.Tensor,
1027	epsilon: float,
1028	quant_dtype: torch.dtype,
1029	) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
1030	y_scale = torch.empty(x.shape[0], 1, dtype=torch.float32, device=x.device)
1031	out = torch.empty(x.shape, dtype=quant_dtype, device=x.device)
1032	residual_out = torch.empty_like(x)
1033
1034	return out, residual_out, y_scale
1035
1036
1037	def _rocm_aiter_rmsnorm_fused_dynamic_quant_impl(
1038	x: torch.Tensor,
1039	weight: torch.Tensor,
1040	epsilon: float,
1041	quant_dtype: torch.dtype,
1042	) -> tuple[torch.Tensor, torch.Tensor]:
1043	import aiter as rocm_aiter
1044
1045	assert quant_dtype in [torch.int8, FP8_DTYPE]
1046
1047	y_scale = torch.empty(x.shape[0], 1, dtype=torch.float32, device=x.device)
1048	out = torch.empty(x.shape, dtype=quant_dtype, device=x.device)
1049
1050	rocm_aiter.rmsnorm2d_fwd_with_dynamicquant(
1051	out, x, y_scale, weight, epsilon, use_model_sensitive_rmsnorm=0
1052	)
1053
1054	return out, y_scale
1055
1056
1057	def _rocm_aiter_rmsnorm_fused_dynamic_quant_fake(
1058	x: torch.Tensor,
1059	weight: torch.Tensor,
1060	epsilon: float,
1061	quant_dtype: torch.dtype,
1062	) -> tuple[torch.Tensor, torch.Tensor]:
1063	y_scale = torch.empty(x.shape[0], 1, dtype=torch.float32, device=x.device)
1064	out = torch.empty(x.shape, dtype=quant_dtype, device=x.device)
1065
1066	return out, y_scale
1067
1068
1069	def _rocm_aiter_fused_allreduce_rmsnorm_impl(
1070	input_: torch.Tensor,
1071	residual: torch.Tensor,
1072	weight: torch.Tensor,
1073	epsilon: float,
1074	gemma_norm: bool = False,
1075	) -> tuple[torch.Tensor, torch.Tensor]:
1076	aiter_ar = rocm_aiter_ops.get_aiter_allreduce()
1077	assert aiter_ar is not None, "aiter allreduce must be initialized"
1078	ca = aiter_ar.aiter_ca
1079	use_1stage = aiter_ar.use_1stage_fused_ar_rms(input_)
1080
1081	if not use_1stage:
1082	from vllm.distributed import get_tp_group
1083
1084	device_comm = get_tp_group().device_communicator
1085	qr_comm = getattr(device_comm, "qr_comm", None)
1086	hidden_dim = input_.shape[-1]
1087	row_size = hidden_dim * input_.element_size()
1088	fused_qr_rmsnorm_ok = (
1089	qr_comm is not None
1090	and not getattr(qr_comm, "disabled", True)
1091	and hasattr(qr_comm, "should_quick_allreduce")
1092	and qr_comm.should_quick_allreduce(input_)
1093	and input_.shape == residual.shape
1094	and input_.dtype in (torch.bfloat16, torch.float16)
1095	and input_.dtype == residual.dtype
1096	and input_.dtype == weight.dtype
1097	and weight.dim() == 1
1098	and weight.numel() == hidden_dim
1099	and input_.numel() % hidden_dim == 0
1100	and row_size > 0
1101	and row_size <= 32 * KB
1102	and (32 * KB) % row_size == 0
1103	)
1104	if fused_qr_rmsnorm_ok:
1105	assert qr_comm is not None
1106	aiter_qr_comm = _get_or_create_aiter_qr_rmsnorm_comm(device_comm)
1107	if aiter_qr_comm is not None:
1108	import aiter
1109
1110	ptr, _, cast_bf2half = aiter_qr_comm
1111	quant_level = qr_comm._get_qr_quant_level(input_)
1112	out = torch.empty_like(input_)
1113	residual_out = torch.empty_like(residual)
1114	aiter.qr_all_reduce_rmsnorm(
1115	ptr,
1116	input_,
1117	residual,
1118	residual_out,
1119	out,
1120	weight,
1121	epsilon,
1122	hidden_dim,
1123	quant_level,
1124	cast_bf2half,
1125	)
1126	return out, residual_out
1127
1128	result = ca.custom_fused_ar_rms(
1129	input_,
1130	residual,
1131	weight,
1132	epsilon,
1133	use_1stage=use_1stage,
1134	gemma_norm=gemma_norm,
1135	)
1136	assert result is not None
1137	return result[0], result[1]
1138
1139
1140	def _rocm_aiter_fused_allreduce_rmsnorm_fake(
1141	input_: torch.Tensor,
1142	residual: torch.Tensor,
1143	weight: torch.Tensor,
1144	epsilon: float,
1145	gemma_norm: bool = False,
1146	) -> tuple[torch.Tensor, torch.Tensor]:
1147	return torch.empty_like(input_), torch.empty_like(residual)
1148
1149
1150	def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_impl(
1151	input_: torch.Tensor,
1152	residual: torch.Tensor,
1153	weight: torch.Tensor,
1154	epsilon: float,
1155	group_size: int,
1156	) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
1157	"""Fused AllReduce + RMSNorm + per-group FP8 quant."""
1158	aiter_ar = rocm_aiter_ops.get_aiter_allreduce()
1159	assert aiter_ar is not None, "aiter allreduce must be initialized"
1160	ca = aiter_ar.aiter_ca
1161
1162	use_1stage = aiter_ar.use_1stage_fused_ar_rms(input_)
1163
1164	result = ca.fused_ar_rms_per_group_quant(
1165	input_,
1166	residual,
1167	w=weight,
1168	eps=epsilon,
1169	group_size=group_size,
1170	registered=torch.cuda.is_current_stream_capturing(),
1171	use_1stage=use_1stage,
1172	)
1173	assert result is not None
1174	return result[0], result[1], result[2]
1175
1176
1177	def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_fake(
1178	input_: torch.Tensor,
1179	residual: torch.Tensor,
1180	weight: torch.Tensor,
1181	epsilon: float,
1182	group_size: int,
1183	) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
1184	hidden_dim = input_.shape[-1]
1185	num_groups = hidden_dim // group_size
1186	quant_out = torch.empty(input_.shape, dtype=FP8_DTYPE, device=input_.device)
1187	residual_out = torch.empty_like(residual)
1188	scale_out = torch.empty(
1189	input_.shape[:-1] + (num_groups,),
1190	dtype=torch.float32,
1191	device=input_.device,
1192	)
1193	return quant_out, residual_out, scale_out
1194
1195
1196	def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_impl(
1197	input_: torch.Tensor,
1198	residual: torch.Tensor,
1199	weight: torch.Tensor,
1200	epsilon: float,
1201	group_size: int,
1202	) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
1203	"""Fused AllReduce + add-RMSNorm + per-group FP8 quant + bf16 normed act.
1204
1205	Wraps the same AITER launcher as ``_rocm_aiter_fused_allreduce_rmsnorm_
1206	quant_per_group_impl`` with ``emit_bf16=True``, emitting the pre-quant
1207	bf16/fp16 normed activation for a parallel consumer (DeepSeek V3.2 sparse
1208	indexer ``wk_weights_proj``).
1209	"""
1210	aiter_ar = rocm_aiter_ops.get_aiter_allreduce()
1211	assert aiter_ar is not None, "aiter allreduce must be initialized"
1212	ca = aiter_ar.aiter_ca
1213
1214	use_1stage = aiter_ar.use_1stage_fused_ar_rms(input_)
1215
1216	result = ca.fused_ar_rms_per_group_quant(
1217	input_,
1218	residual,
1219	w=weight,
1220	eps=epsilon,
1221	group_size=group_size,
1222	registered=torch.cuda.is_current_stream_capturing(),
1223	use_1stage=use_1stage,
1224	emit_bf16=True,
1225	)
1226	assert result is not None
1227	assert len(result) == 4, "emit_bf16=True must return four tensors from aiter"
1228	return result[0], result[1], result[2], result[3]
1229
1230
1231	def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_fake(
1232	input_: torch.Tensor,
1233	residual: torch.Tensor,
1234	weight: torch.Tensor,
1235	epsilon: float,
1236	group_size: int,
1237	) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
1238	hidden_dim = input_.shape[-1]
1239	num_groups = hidden_dim // group_size
1240	quant_out = torch.empty(input_.shape, dtype=FP8_DTYPE, device=input_.device)
1241	residual_out = torch.empty_like(residual)
1242	scale_out = torch.empty(
1243	input_.shape[:-1] + (num_groups,),
1244	dtype=torch.float32,
1245	device=input_.device,
1246	)
1247	bf16_norm_out = torch.empty_like(input_)
1248	return quant_out, residual_out, scale_out, bf16_norm_out
1249
1250
1251	def _rocm_aiter_per_tensor_quant_impl(
1252	out: torch.Tensor,
1253	x: torch.Tensor,
1254	scale: torch.Tensor,
1255	is_dynamic: bool,
1256	) -> None:
1257	from aiter.ops.quant import dynamic_per_tensor_quant, static_per_tensor_quant
1258
1259	if is_dynamic:
1260	dynamic_per_tensor_quant(out, x, scale)
1261	else:
1262	static_per_tensor_quant(out, x, scale)
1263
1264
1265	def _rocm_aiter_per_token_quant_impl(
1266	x: torch.Tensor, quant_dtype: torch.dtype, scale: torch.Tensor \| None = None
1267	) -> tuple[torch.Tensor, torch.Tensor]:
1268	from aiter.ops.quant import dynamic_per_token_scaled_quant
1269
1270	assert quant_dtype in [torch.int8, FP8_DTYPE]
1271
1272	out_shape = x.shape
1273	out = torch.empty(x.shape, dtype=quant_dtype, device=x.device)
1274	if scale is None:
1275	scale = torch.empty((*out_shape[:-1], 1), dtype=torch.float32, device=x.device)
1276	dynamic_per_token_scaled_quant(
1277	out,
1278	x,
1279	scale,
1280	scale_ub=None,
1281	shuffle_scale=False,
1282	num_rows=None,
1283	num_rows_factor=1,
1284	)
1285	return out, scale
1286
1287
1288	def _rocm_aiter_per_token_quant_fake(
1289	x: torch.Tensor, quant_dtype: torch.dtype, scale: torch.Tensor \| None = None
1290	) -> tuple[torch.Tensor, torch.Tensor]:
1291	out_shape = x.shape
1292	return (
1293	torch.empty(x.shape, dtype=quant_dtype, device=x.device),
1294	torch.empty((*out_shape[:-1], 1), dtype=torch.float32, device=x.device),
1295	)
1296
1297
1298	def _rocm_aiter_rmsnorm_with_add_fp8_group_quant_impl(
1299	x: torch.Tensor,
1300	residual: torch.Tensor,
1301	weight: torch.Tensor,
1302	variance_epsilon: float,
1303	group_size: int,
1304	) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
1305	from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant
1306
1307	(x_quant, x_quant_scales), _, _, res = fused_rms_fp8_group_quant(
1308	x,
1309	weight,
1310	variance_epsilon,
1311	None,
1312	None,
1313	None,
1314	group_size=group_size,
1315	dtype_quant=FP8_DTYPE,
1316	res1=residual,
1317	)
1318	return (
1319	x_quant,
1320	res,
1321	x_quant_scales,
1322	)
1323
1324
1325	def _rocm_aiter_rmsnorm_with_add_fp8_group_quant_fake(
1326	x: torch.Tensor,
1327	residual: torch.Tensor,
1328	weight: torch.Tensor,
1329	variance_epsilon: float,
1330	group_size: int,
1331	) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
1332	M, N = x.shape
1333	scale_shape = (M, (N + group_size - 1) // group_size)
1334	return (
1335	torch.empty_like(x, dtype=FP8_DTYPE, device=x.device),
1336	torch.empty_like(residual, device=residual.device),
1337	torch.empty(scale_shape, dtype=torch.float32, device=x.device),
1338	)
1339
1340
1341	def _rocm_aiter_rmsnorm_fp8_group_quant_impl(
1342	x: torch.Tensor,
1343	weight: torch.Tensor,
1344	variance_epsilon: float,
1345	group_size: int,
1346	) -> tuple[torch.Tensor, torch.Tensor]:
1347	from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant
1348
1349	(x_quant, x_quant_scales), _, _, res = fused_rms_fp8_group_quant(
1350	x,
1351	weight,
1352	variance_epsilon,
1353	None,
1354	None,
1355	None,
1356	group_size=group_size,
1357	dtype_quant=FP8_DTYPE,
1358	res1=None,
1359	)
1360	return (x_quant, x_quant_scales)
1361
1362
1363	def _rocm_aiter_rmsnorm_fp8_group_quant_fake(
1364	x: torch.Tensor,
1365	weight: torch.Tensor,
1366	variance_epsilon: float,
1367	group_size: int,
1368	) -> tuple[torch.Tensor, torch.Tensor]:
1369	M, N = x.shape
1370	scale_shape = (M, (N + group_size - 1) // group_size)
1371	return (
1372	torch.empty_like(x, dtype=FP8_DTYPE, device=x.device),
1373	torch.empty(scale_shape, dtype=torch.float32, device=x.device),
1374	)
1375
1376
1377	def _rocm_aiter_fused_rms_gated_fp8_group_quant_impl(
1378	x: torch.Tensor,
1379	weight: torch.Tensor,
1380	bias: torch.Tensor \| None,
1381	z: torch.Tensor,
1382	eps: float,
1383	norm_before_gate: bool,
1384	activation: str,
1385	group_size: int,
1386	) -> tuple[torch.Tensor, torch.Tensor]:
1387	"""Fused gated-RMSNorm + FP8 group quantization via aiter Triton kernel."""
1388	from aiter.ops.triton.quant import fused_rms_gated_fp8_group_quant
1389
1390	return fused_rms_gated_fp8_group_quant(
1391	x,
1392	weight,
1393	bias,
1394	z,
1395	eps,
1396	norm_before_gate=norm_before_gate,
1397	activation=activation,
1398	out_dtype=FP8_DTYPE,
1399	group_size=group_size,
1400	)
1401
1402
1403	def _rocm_aiter_fused_rms_gated_fp8_group_quant_fake(
1404	x: torch.Tensor,
1405	weight: torch.Tensor,
1406	bias: torch.Tensor \| None,
1407	z: torch.Tensor,
1408	eps: float,
1409	norm_before_gate: bool,
1410	activation: str,
1411	group_size: int,
1412	) -> tuple[torch.Tensor, torch.Tensor]:
1413	M, N = x.shape
1414	scale_shape = (M, (N + group_size - 1) // group_size)
1415	return (
1416	torch.empty_like(x, dtype=FP8_DTYPE, device=x.device),
1417	torch.empty(scale_shape, dtype=torch.float32, device=x.device),
1418	)
1419
1420
1421	def _rocm_aiter_group_fp8_quant_impl(
1422	x: torch.Tensor,
1423	group_size: int,
1424	transpose_scale: bool = False,
1425	) -> tuple[torch.Tensor, torch.Tensor]:
1426	assert x.shape[-1] % group_size == 0, "Input shape must be divisible by group size"
1427	from aiter import QuantType, get_hip_quant
1428
1429	aiter_per1x128_quant = get_hip_quant(QuantType.per_1x128)
1430	return aiter_per1x128_quant(
1431	x.contiguous(), quant_dtype=FP8_DTYPE, transpose_scale=transpose_scale
1432	)
1433
1434
1435	def _rocm_aiter_group_fp8_quant_fake(
1436	x: torch.Tensor,
1437	group_size: int,
1438	transpose_scale: bool = False,
1439	) -> tuple[torch.Tensor, torch.Tensor]:
1440	M, N = x.shape
1441	x_fp8 = torch.empty((M, N), dtype=FP8_DTYPE, device=x.device)
1442	out_bs = torch.empty(
1443	(
1444	M,
1445	(N + group_size - 1) // group_size,
1446	),
1447	dtype=torch.float32,
1448	device=x.device,
1449	)
1450	return x_fp8, out_bs
1451
1452
1453	def _rocm_aiter_act_mul_and_fp8_group_quant_impl(
1454	x: torch.Tensor,
1455	group_size: int,
1456	) -> tuple[torch.Tensor, torch.Tensor]:
1457	from aiter.ops.triton.activation import act_mul_and_fp8_group_quant
1458
1459	return act_mul_and_fp8_group_quant(
1460	x,
1461	activation="silu",
1462	group_size=group_size,
1463	dtype_quant=FP8_DTYPE,
1464	)
1465
1466
1467	def _rocm_aiter_act_mul_and_fp8_group_quant_fake(
1468	x: torch.Tensor,
1469	group_size: int,
1470	) -> tuple[torch.Tensor, torch.Tensor]:
1471	M, N = x.shape
1472	assert N % 2 == 0
1473	N_half = N // 2
1474	x_fp8 = torch.empty((M, N_half), dtype=FP8_DTYPE, device=x.device)
1475	out_bs = torch.empty(
1476	(
1477	M,
1478	(N_half + group_size - 1) // group_size,
1479	),
1480	dtype=torch.float32,
1481	device=x.device,
1482	)
1483	return x_fp8, out_bs
1484
1485
1486	def _rocm_aiter_triton_add_rmsnorm_pad_impl(
1487	x: torch.Tensor,
1488	weight: torch.Tensor,
1489	variance_epsilon: float,
1490	residual: torch.Tensor,
1491	x_pad_to_multiple: int,
1492	) -> tuple[torch.Tensor, torch.Tensor]:
1493	from aiter.ops.triton.fused_add_rmsnorm_pad import fused_add_rmsnorm_pad
1494
1495	return fused_add_rmsnorm_pad(
1496	x,
1497	weight,
1498	variance_epsilon,
1499	residual,
1500	x_pad_to_multiple=x_pad_to_multiple,
1501	)
1502
1503
1504	def _rocm_aiter_triton_add_rmsnorm_pad_fake(
1505	x: torch.Tensor,
1506	weight: torch.Tensor,
1507	variance_epsilon: float,
1508	residual: torch.Tensor,
1509	x_pad_to_multiple: int,
1510	) -> tuple[torch.Tensor, torch.Tensor]:
1511	M, N = x.shape
1512	if x_pad_to_multiple > 0:
1513	N_out = (N + x_pad_to_multiple - 1) // x_pad_to_multiple * x_pad_to_multiple
1514	else:
1515	N_out = N
1516	out = torch.empty((M, N_out), dtype=x.dtype, device=x.device)
1517	residual_out = torch.empty_like(residual)
1518	return out, residual_out
1519
1520
1521	def _fused_mla_dual_rms_norm_impl(
1522	x1: torch.Tensor,
1523	x1_weight: torch.Tensor,
1524	x2: torch.Tensor,
1525	x2_weight: torch.Tensor,
1526	x1_epsilon: float,
1527	x2_epsilon: float,
1528	) -> tuple[torch.Tensor, torch.Tensor]:
1529	from aiter.ops.fused_qk_norm_rope_cache_quant import _fused_qk_rmsnorm
1530
1531	return _fused_qk_rmsnorm(
1532	q_out=None,
1533	q=x1,
1534	q_weight=x1_weight,
1535	q_eps=x1_epsilon,
1536	k_out=None,
1537	k=x2,
1538	k_weight=x2_weight,
1539	k_eps=x2_epsilon,
1540	)
1541
1542
1543	def _fused_mla_dual_rms_norm_fake(
1544	x1: torch.Tensor,
1545	x1_weight: torch.Tensor,
1546	x2: torch.Tensor,
1547	x2_weight: torch.Tensor,
1548	x1_epsilon: float,
1549	x2_epsilon: float,
1550	) -> tuple[torch.Tensor, torch.Tensor]:
1551	return (torch.empty_like(x1), torch.empty_like(x2))
1552
1553
1554	def _fused_mla_dual_rms_norm_per_token_quant_impl(
1555	q: torch.Tensor,
1556	q_weight: torch.Tensor,
1557	kv: torch.Tensor,
1558	kv_weight: torch.Tensor,
1559	q_epsilon: float,
1560	kv_epsilon: float,
1561	) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
1562	"""Fused MLA q/kv RMSNorm (+ FP8 per-token quant on q) via AITER.
1563
1564	Backs the ``fused_mla_dual_rms_norm_per_token_quant`` custom op used by the
1565	MLA FP8 attention fusion when the q latent is quantized per token (a single
1566	``(M, 1)`` scale). Only the q latent is FP8 quantized (it feeds the
1567	FP8 ``q_b_proj`` GEMM); the kv latent is RMS-normed and consumed by attention as bf16.
1568	"""
1569	from aiter.ops.fused_qk_rmsnorm_group_quant import (
1570	fused_qk_rmsnorm_per_token_quant,
1571	)
1572
1573	mq, nq = q.shape
1574	q_out = torch.empty((mq, nq), dtype=FP8_DTYPE, device=q.device)
1575	q_scale = torch.empty((mq, 1), dtype=torch.float32, device=q.device)
1576	kv_normed = torch.empty(kv.shape, dtype=kv.dtype, device=kv.device)
1577
1578	# q -> RMSNorm + FP8 per-token quant (q slot); kv -> RMSNorm only (k slot).
1579	# `split` views are accepted directly (unit inner stride); the kernel
1580	# handles strided inputs, matching the aiter op-test usage.
1581	fused_qk_rmsnorm_per_token_quant(
1582	q_out_quantized=q_out,
1583	q_out_scale=q_scale,
1584	q=q,
1585	q_weight=q_weight,
1586	q_epsilon=q_epsilon,
1587	k_out=kv_normed,
1588	k=kv,
1589	k_weight=kv_weight,
1590	k_epsilon=kv_epsilon,
1591	gemma_norm=False,
1592	)
1593	return q_out, q_scale, kv_normed
1594
1595
1596	def _fused_mla_dual_rms_norm_per_token_quant_fake(
1597	q: torch.Tensor,
1598	q_weight: torch.Tensor,
1599	kv: torch.Tensor,
1600	kv_weight: torch.Tensor,
1601	q_epsilon: float,
1602	kv_epsilon: float,
1603	) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
1604	mq, nq = q.shape
1605	q_out = torch.empty((mq, nq), dtype=FP8_DTYPE, device=q.device)
1606	q_scale = torch.empty((mq, 1), dtype=torch.float32, device=q.device)
1607	kv_normed = torch.empty(kv.shape, dtype=kv.dtype, device=kv.device)
1608	return q_out, q_scale, kv_normed
1609
1610
1611	def _fused_mla_dual_rms_norm_group_quant_impl(
1612	q: torch.Tensor,
1613	q_weight: torch.Tensor,
1614	kv: torch.Tensor,
1615	kv_weight: torch.Tensor,
1616	q_epsilon: float,
1617	kv_epsilon: float,
1618	group_size: int,
1619	transpose_scale: bool,
1620	) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
1621	"""Fused MLA q/kv RMSNorm (+ FP8 per-1x128 group quant on q) via AITER.
1622
1623	Backs the ``fused_mla_dual_rms_norm_group_quant`` custom op used by the
1624	DeepSeek-V4 ROCm attention path when the q latent is group-quantized
1625	(``(M, N/128)`` scales) for the block-scaled FP8 wq_b GEMMs. The q latent
1626	is quantized directly from the fp32 norm accumulator, skipping the bf16
1627	intermediate the standalone quant path would re-read; the kv latent is
1628	RMS-normed to bf16 and consumed by the fused insert kernel unchanged.
1629	"""
1630	from aiter.ops.fused_qk_rmsnorm_group_quant import (
1631	fused_qk_rmsnorm_group_quant,
1632	)
1633
1634	mq, nq = q.shape
1635	q_out = torch.empty((mq, nq), dtype=FP8_DTYPE, device=q.device)
1636	q_scale = torch.empty((mq, nq // group_size), dtype=torch.float32, device=q.device)
1637	kv_normed = torch.empty(kv.shape, dtype=kv.dtype, device=kv.device)
1638
1639	# q -> RMSNorm + FP8 group quant (q slot); kv -> RMSNorm only (k slot).
1640	# `split` views are accepted directly (unit inner stride); the kernel
1641	# handles strided inputs, matching the aiter op-test usage.
1642	fused_qk_rmsnorm_group_quant(
1643	q_out_quantized=q_out,
1644	q_out_scale=q_scale,
1645	q=q,
1646	q_weight=q_weight,
1647	q_epsilon=q_epsilon,
1648	k_out=kv_normed,
1649	k=kv,
1650	k_weight=kv_weight,
1651	k_epsilon=kv_epsilon,
1652	group_size=group_size,
1653	transpose_scale=transpose_scale,
1654	)
1655	return q_out, q_scale, kv_normed
1656
1657
1658	def _fused_mla_dual_rms_norm_group_quant_fake(
1659	q: torch.Tensor,
1660	q_weight: torch.Tensor,
1661	kv: torch.Tensor,
1662	kv_weight: torch.Tensor,
1663	q_epsilon: float,
1664	kv_epsilon: float,
1665	group_size: int,
1666	transpose_scale: bool,
1667	) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
1668	mq, nq = q.shape
1669	q_out = torch.empty((mq, nq), dtype=FP8_DTYPE, device=q.device)
1670	q_scale = torch.empty((mq, nq // group_size), dtype=torch.float32, device=q.device)
1671	kv_normed = torch.empty(kv.shape, dtype=kv.dtype, device=kv.device)
1672	return q_out, q_scale, kv_normed
1673
1674
1675	def _rocm_aiter_gemm_a8wfp4_impl(
1676	x: torch.Tensor,
1677	w: torch.Tensor,
1678	x_scales: torch.Tensor,
1679	w_scales: torch.Tensor,
1680	out_dtype: torch.dtype,
1681	) -> torch.Tensor:
1682	from aiter.ops.triton.gemm_a8wfp4 import gemm_a8wfp4
1683
1684	M, N = x.shape[0], w.shape[0]
1685	y = torch.empty(M, N, dtype=out_dtype, device=x.device)
1686	gemm_a8wfp4(
1687	x=x,
1688	w=w,
1689	y=y,
1690	x_scales=x_scales,
1691	w_scales=w_scales,
1692	dtype=out_dtype,
1693	config=None,
1694	)
1695	return y
1696
1697
1698	def _rocm_aiter_gemm_a8wfp4_fake(
1699	x: torch.Tensor,
1700	w: torch.Tensor,
1701	x_scales: torch.Tensor,
1702	w_scales: torch.Tensor,
1703	out_dtype: torch.dtype,
1704	) -> torch.Tensor:
1705	return torch.empty(x.shape[0], w.shape[0], dtype=out_dtype, device=x.device)
1706
1707
1708	def _triton_rotary_embedding_impl(
1709	positions: torch.Tensor,
1710	query: torch.Tensor,
1711	key: torch.Tensor,
1712	head_size: int,
1713	cos_sin_cache: torch.Tensor,
1714	is_neox: bool,
1715	offsets: torch.Tensor \| None = None,
1716	) -> None:
1717	# Modifies query and key in-place
1718	from aiter.ops.triton.rope.rope import (
1719	rope_cached_thd_positions_offsets_2c_fwd_inplace,
1720	)
1721
1722	num_tokens = positions.numel()
1723	cos, sin = cos_sin_cache.chunk(2, dim=-1)
1724	query_shape = query.shape
1725	key_shape = key.shape
1726	rotate_style = 0 if is_neox else 1
1727	rotary_dim = head_size
1728
1729	query = query.view(num_tokens, -1, head_size)
1730	key = key.view(num_tokens, -1, head_size)
1731	query_ = query[..., :rotary_dim]
1732	key_ = key[..., :rotary_dim]
1733	positions = positions.view(*query.shape[:1])
1734	rope_cached_thd_positions_offsets_2c_fwd_inplace(
1735	query_,
1736	key_,
1737	cos,
1738	sin,
1739	positions,
1740	offsets,
1741	rotate_style,
1742	reuse_freqs_front_part=True,
1743	nope_first=False,
1744	)
1745	query = query.view(query_shape)
1746	key = key.view(key_shape)
1747
1748
1749	def _rocm_aiter_fp8_attn_impl(
1750	q: torch.Tensor,
1751	k: torch.Tensor,
1752	v: torch.Tensor,
1753	q_descale: torch.Tensor,
1754	k_descale: torch.Tensor,
1755	v_descale: torch.Tensor,
1756	batch_size: int,
1757	output_dtype: torch.dtype,
1758	scale: float \| None = None,
1759	cu_seqlens: torch.Tensor \| None = None,
1760	max_seqlen: torch.Tensor \| None = None,
1761	) -> torch.Tensor:
1762	"""Run AITER FP8 attention for fixed or packed inputs."""
1763	from aiter import flash_attn_varlen_fp8_pertensor_func
1764
1765	q_len = q.size(1)
1766	if cu_seqlens is None:
1767	cu_seqlens = torch.arange(
1768	0, (batch_size + 1) * q_len, step=q_len, dtype=torch.int32, device=q.device
1769	)
1770	max_seqlen_value = q_len if max_seqlen is None else max_seqlen.item()
1771
1772	q, k, v = (x.flatten(0, 1) for x in (q, k, v))
1773	output = flash_attn_varlen_fp8_pertensor_func(
1774	q,
1775	k,
1776	v,
1777	q_descale=q_descale,
1778	k_descale=k_descale,
1779	v_descale=v_descale,
1780	cu_seqlens_q=cu_seqlens,
1781	cu_seqlens_k=cu_seqlens,
1782	max_seqlen_q=max_seqlen_value,
1783	max_seqlen_k=max_seqlen_value,
1784	causal=False,
1785	softmax_scale=scale,
1786	)
1787	return output.to(output_dtype).reshape(batch_size, q_len, *output.shape[1:])
1788
1789
1790	def _rocm_aiter_fp8_attn_fake(
1791	q: torch.Tensor,
1792	k: torch.Tensor,
1793	v: torch.Tensor,
1794	q_descale: torch.Tensor,
1795	k_descale: torch.Tensor,
1796	v_descale: torch.Tensor,
1797	batch_size: int,
1798	output_dtype: torch.dtype,
1799	scale: float \| None = None,
1800	cu_seqlens: torch.Tensor \| None = None,
1801	max_seqlen: torch.Tensor \| None = None,
1802	) -> torch.Tensor:
1803	return torch.empty(
1804	(*q.shape[:-1], v.shape[-1]), device=q.device, dtype=output_dtype
1805	)
1806
1807
1808	def _mhc_delayed_pre_tail(
1809	gemm_out: torch.Tensor,
1810	sqrsum: torch.Tensor,
1811	residual: torch.Tensor,
1812	hc_scale: torch.Tensor,
1813	hc_base: torch.Tensor,
1814	pre_mix: torch.Tensor \| None,
1815	rms_eps: float,
1816	hc_pre_eps: float,
1817	hc_sinkhorn_eps: float,
1818	hc_post_mult_value: float,
1819	sinkhorn_repeat: int,
1820	) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
1821	"""Finish a delayed mHC pre once AITER has produced the projection.
1822
1823	Shared by the plain and the post-fused entry points, which differ only in
1824	which kernel computes ``gemm_out`` / ``sqrsum`` and the residual they
1825	project. ``mhc_pre_big_fuse`` also writes a collapse against its own
1826	pre-mix, which the delayed formulation cannot use and so discards.
1827	Returns ``(post_mix, comb_mix, layer_input, next_pre_mix)``.
1828	"""
1829	from aiter.ops.mhc import mhc_pre_big_fuse
1830
1831	num_tokens, hc_mult, hidden_size = residual.shape
1832	device = residual.device
1833	with torch.device(device):
1834	post_mix = torch.empty(
1835	num_tokens, hc_mult, 1, dtype=torch.float32, device=device
1836	)
1837	comb_mix = torch.empty(
1838	num_tokens, hc_mult, hc_mult, dtype=torch.float32, device=device
1839	)
1840	unused_collapse = torch.empty(
1841	num_tokens, hidden_size, dtype=torch.bfloat16, device=device
1842	)
1843	mhc_pre_big_fuse(
1844	post_mix,
1845	comb_mix,
1846	unused_collapse,
1847	gemm_out,
1848	sqrsum,
1849	hc_scale,
1850	hc_base,
1851	residual,
1852	rms_eps,
1853	hc_pre_eps,
1854	hc_sinkhorn_eps,
1855	hc_post_mult_value,
1856	sinkhorn_repeat,
1857	)
1858
1859	next_pre_mix = torch.ops.vllm.mhc_pre_mix_triton(
1860	gemm_out,
1861	sqrsum,
1862	hc_scale,
1863	hc_base,
1864	hc_mult,
1865	hc_mult * hidden_size,
1866	rms_eps,
1867	hc_pre_eps,
1868	)
1869	if pre_mix is None:
1870	# Model entry selects residual stream zero.
1871	layer_input = residual[:, 0]
1872	else:
1873	layer_input = torch.ops.vllm.hc_collapse_triton(residual, pre_mix)
1874	return post_mix, comb_mix, layer_input, next_pre_mix
1875
1876
1877	# Global flag to ensure ops are registered only once
1878	_OPS_REGISTERED = False
1879
1880
1881	_AITER_SITUV2_ACT_ENV = {
1882	"a8w4": "AITER_SITUV2_A8W4",
1883	"a4w4": "AITER_SITUV2_A4W4",
1884	}
1885
1886
1887	def _resolve_situv2_activation() -> str:
1888	"""VLLM_ROCM_USE_AITER_MOE_SITUV2 -> a4w4 \| a8w4 \| a16w4.
1889
1890	auto (default) and the legacy 1 mean a4w4; legacy 0 means a16w4.
1891	"""
1892	value = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2.lower()
1893	if value in ("auto", "1"):
1894	return "a4w4"
1895	if value == "0":
1896	return "a16w4"
1897	return value
1898
1899
1900	def _sync_aiter_situv2_moe_env() -> None:
1901	"""Set the AITER_SITUV2_* env matching VLLM_ROCM_USE_AITER_MOE_SITUV2.
1902
1903	AITER reads AITER_SITUV2_A8W4 / AITER_SITUV2_A4W4 and defaults to a16w4,
1904	so exactly one is set for a8w4/a4w4 and both are cleared for a16w4.
1905	"""
1906	import os
1907
1908	selected = _AITER_SITUV2_ACT_ENV.get(_resolve_situv2_activation())
1909	for name in _AITER_SITUV2_ACT_ENV.values():
1910	if name == selected:
1911	os.environ[name] = "1"
1912	else:
1913	os.environ.pop(name, None)
1914
1915
1916	_GFX950_C4A_AITER_MAX_COMPRESSED_SEQ_LEN = 64 * 1024
1917	_GFX950_C4A_NATIVE_MAX_ROWS = 256
1918	_GFX950_DSV4_NATIVE_MAX_COLUMNS = 1024 * 1024
1919
1920
1921	class rocm_aiter_ops:
1922	"""ROCm AITER operations wrapper for AMD GPU acceleration in vLLM.
1923
1924	This class centralizes the import and registration of AITER ops,
1925	and provides a unified interface for checking if AITER is enabled.
1926	Operations are only available on supported gfx9
1927	architectures when aiter is installed.
1928
1929	The class uses environment variables to control which features are enabled,
1930	allowing fine-grained control over which AITER optimizations are used.
1931
1932	Environment Variables:
1933	VLLM_ROCM_USE_AITER: Main toggle for all AITER operations.
1934	VLLM_ROCM_USE_AITER_LINEAR: Controls GEMM and quantization ops.
1935	VLLM_ROCM_USE_AITER_RMSNORM: Controls RMSNorm operations.
1936	VLLM_ROCM_USE_AITER_MOE: Controls MoE (Mixture of Experts) ops.
1937	VLLM_ROCM_USE_AITER_MLA: Controls MLA (Multi-head Latent Attention) ops.
1938	VLLM_ROCM_USE_AITER_MHA: Controls MHA ops including flash_attn_varlen.
1939	VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION: Controls Triton unified attention.
1940	VLLM_ROCM_USE_AITER_FP8BMM: Controls FP8 batched matrix multiply.
1941	VLLM_ROCM_USE_AITER_FP4_ASM_GEMM: Controls FP4 assembly GEMM.
1942	VLLM_ROCM_USE_AITER_TRITON_ROPE: Controls Triton rotary embeddings.
1943	VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA: Controls Triton sparse MLA (gfx950).
1944	VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: Controls shared expert fusion.
1945	VLLM_ROCM_USE_AITER_MOE_SITUV2: SiTUv2 FlyDSL MoE activation
1946	dtype (a16w4 \| a8w4 \| a4w4).
1947	VLLM_ROCM_USE_AITER_TRITON_GEMM: Controls Triton unquantized GEMM.
1948
1949	Note:
1950	The environment variables are assigned when the module is imported,
1951	so you can't change the environment variables after the module is imported.
1952	This is done out of performance consideration. Accessing environment variables
1953	is expensive as described in issue https://github.com/vllm-project/vllm/issues/17067
1954	so we don't want to do it repeatedly, especially in the hot path (the forward pass).
1955	You can call the refresh_env_variables() function to reload the env variables
1956	after monkey patching the env variables in the unit test.
1957
1958	Check Functions:
1959	All check functions (is_*_enabled) are decorated with @if_aiter_supported,
1960	which verifies: (1) platform is ROCm, (2) device arch is gfx9, and
1961	(3) aiter library is installed. The check function then also verifies
1962	the corresponding environment variable is enabled.
1963	i.e. ___
1964	is_enabled() == current_platform.is_rocm() and \| checked by
1965	current_platform.is_on_gfx9() and \| @if_aiter_supported
1966	IS_AITER_FOUND and _______________\|
1967	cls._AITER_ENABLED -----> Check by the logic in `is_enabled()`
1968
1969	Example:
1970	from vllm._aiter_ops import rocm_aiter_ops
1971
1972	# Check if aiter is enabled before using operations
1973	if rocm_aiter_ops.is_enabled():
1974	result = rocm_aiter_ops.per_token_quant(x, FP8_DTYPE)
1975
1976	Operations:
1977	- GEMM operations: gemm_a8w8, gemm_a8w8_blockscale
1978	- Fused MoE: fused_moe, asm_moe_tkw1
1979	- Routing: topk_softmax, biased_grouped_topk, grouped_topk
1980	- MLA decode: mla_decode_fwd
1981	- Quantization: per_tensor_quant, per_token_quant, group_fp8_quant
1982	- Triton ops: triton_rotary_embed, triton_fp8_bmm, triton_gemm_a8w8_blockscale
1983
1984	"""
1985
1986	_MOE_DISPATCH_POLICY: int \| None = None
1987
1988	@classmethod
1989	@if_aiter_supported
1990	def get_moe_dispatch_policy(cls) -> int:
1991	"""Cached MoE sorting dispatch policy."""
1992	if cls._MOE_DISPATCH_POLICY is None:
1993	import vllm.envs as envs
1994
1995	cls._MOE_DISPATCH_POLICY = envs.VLLM_ROCM_AITER_MOE_DISPATCH_POLICY
1996	return cls._MOE_DISPATCH_POLICY
1997
1998	# Check if the env variable is set
1999	_AITER_ENABLED = envs.VLLM_ROCM_USE_AITER
2000	_CUSTOM_ALL_REDUCE_ENABLED = envs.VLLM_ROCM_USE_AITER_CUSTOM_AR
2001	_LINEAR_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR
2002	_FMOE_ENABLED = envs.VLLM_ROCM_USE_AITER_MOE
2003	_MLA_ENABLED = envs.VLLM_ROCM_USE_AITER_MLA
2004	_MHA_ENABLED = envs.VLLM_ROCM_USE_AITER_MHA
2005	_SHUFFLE_KV_CACHE_ENABLED = envs.VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT
2006	_TRITON_UNIFIED_ATTN_ENABLED = envs.VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION
2007	# TODO: Consolidate under _LINEAR_ENABLED
2008	_FP8BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP8BMM
2009	_FP4BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP4BMM
2010	_LINEAR_HIPBMM_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM
2011	# TODO: Consolidate under _LINEAR_ENABLED
2012	_FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM
2013	# TODO: Consolidate under VLLM_ROCM_USE_AITER_ROPE
2014	_TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE
2015	_TRITON_SPARSE_MLA = envs.VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA
2016	_MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
2017	_MOE_SITUV2 = _resolve_situv2_activation()
2018	# TODO: Consolidate under _LINEAR_ENABLED
2019	_TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM
2020	# Lazily probed: whether aiter.topk_softmax supports the
2021	# num_shared_experts / shared_expert_scoring_func args (7-arg form).
2022	_TOPK_SOFTMAX_FUSED_SIGMOID: bool \| None = None
2023
2024	@classmethod
2025	def refresh_env_variables(cls):
2026	"""Since the environment variables are assigned when the module is imported,
2027	This is a helper function to reload all the env variables from
2028	the environment variables.
2029	for example, after monkey patching the env variables in the unit test,
2030	you can call this function to reload the env variables.
2031	"""
2032	cls._AITER_ENABLED = envs.VLLM_ROCM_USE_AITER
2033	cls._CUSTOM_ALL_REDUCE_ENABLED = envs.VLLM_ROCM_USE_AITER_CUSTOM_AR
2034	cls._LINEAR_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR
2035	cls._FMOE_ENABLED = envs.VLLM_ROCM_USE_AITER_MOE
2036	cls._MLA_ENABLED = envs.VLLM_ROCM_USE_AITER_MLA
2037	cls._MHA_ENABLED = envs.VLLM_ROCM_USE_AITER_MHA
2038	cls._SHUFFLE_KV_CACHE_ENABLED = envs.VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT
2039	cls._TRITON_UNIFIED_ATTN_ENABLED = envs.VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION
2040	cls._FP8BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP8BMM
2041	cls._FP4BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP4BMM
2042	cls._LINEAR_HIPBMM_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM
2043	cls._FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM
2044	cls._TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE
2045	cls._TRITON_SPARSE_MLA = envs.VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA
2046	cls._MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
2047	cls._MOE_SITUV2 = _resolve_situv2_activation()
2048	_sync_aiter_situv2_moe_env()
2049	cls._TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM
2050	cls._MOE_DISPATCH_POLICY = envs.VLLM_ROCM_AITER_MOE_DISPATCH_POLICY
2051
2052	@staticmethod
2053	def get_aiter_activation_type(activation_str: str) -> "ActivationType \| None":
2054	"""Given an activation type as a string, returns the corresponding aiter ActivationType enum.
2055	Supported activation types: "no", "none", "silu", "gelu", "swiglu".
2056	Returns None if the mapping fails.
2057
2058	Args:
2059	activation_str (str): Activation type as string.
2060
2061	Returns:
2062	Aiter ActivationType enum value, or None if not found.
2063
2064	"""
2065	# Import only locally, since aiter may not always be available.
2066	try:
2067	from aiter import ActivationType
2068	except ImportError:
2069	return None
2070
2071	if not isinstance(activation_str, str):
2072	return None
2073
2074	name = activation_str.strip().lower()
2075	mapping = {
2076	"none": ActivationType.No,
2077	"no": ActivationType.No,
2078	"silu": ActivationType.Silu,
2079	"gelu": ActivationType.Gelu,
2080	"swiglu": ActivationType.Swiglu,
2081	"situ": getattr(ActivationType, "Situv2", None),
2082	}
2083	return mapping.get(name)
2084
2085	@staticmethod
2086	def get_aiter_quant_type(quant_type_str: str) -> "QuantType \| None":
2087	"""Given a quantization type as a string, returns the corresponding aiter QuantType enum.
2088	Supported quantization types: "no", "per_tensor", "per_token", "per_1x32", "per_1x128", "per_128x128".
2089	Returns None if the mapping fails.
2090
2091	Args:
2092	quant_type_str (str): Quantization type as string.
2093
2094	Returns:
2095	Aiter QuantType enum value, or None if not found.
2096
2097	"""
2098	try:
2099	from aiter import QuantType
2100	except ImportError:
2101	return None
2102
2103	if not isinstance(quant_type_str, str):
2104	return None
2105
2106	name = quant_type_str.strip().lower()
2107	mapping = {
2108	"no": QuantType.No,
2109	"per_tensor": QuantType.per_Tensor,
2110	"per_token": QuantType.per_Token,
2111	"per_1x32": QuantType.per_1x32,
2112	"per_1x128": QuantType.per_1x128,
2113	"per_128x128": QuantType.per_128x128,
2114	}
2115	return mapping.get(name)
2116
2117	@classmethod
2118	@if_aiter_supported
2119	def is_enabled(cls) -> bool:
2120	return cls._AITER_ENABLED
2121
2122	@classmethod
2123	def is_rdna_aiter_enabled(cls) -> bool:
2124	"""AITER on RDNA4 (gfx12): library present, arch is rdna4, and the user
2125	enabled aiter. Only aiter's Triton kernels exist on gfx12 (no CK build),
2126	so this deliberately stays off the gfx9/CK `@if_aiter_supported` umbrella
2127	and gates only the Triton paths rdna4 uses. The gfx12 analog of
2128	`is_enabled()`."""
2129	if not current_platform.is_rocm() or not IS_AITER_FOUND:
2130	return False
2131	from vllm.platforms.rocm import on_rdna4
2132
2133	return on_rdna4() and cls._AITER_ENABLED
2134
2135	@classmethod
2136	def is_rdna_linear_enabled(cls) -> bool:
2137	"""RDNA4 (gfx12) analog of is_linear_enabled() (aiter Triton blockscale)."""
2138	return cls.is_rdna_aiter_enabled() and cls._LINEAR_ENABLED
2139
2140	@classmethod
2141	@if_aiter_supported
2142	def is_linear_enabled(cls) -> bool:
2143	return cls._AITER_ENABLED and cls._LINEAR_ENABLED
2144
2145	@classmethod
2146	@if_aiter_supported
2147	def is_linear_fp8_enabled(cls) -> bool:
2148	return cls.is_linear_enabled()
2149
2150	@classmethod
2151	@if_aiter_supported
2152	def is_fused_moe_enabled(cls) -> bool:
2153	return cls._AITER_ENABLED and cls._FMOE_ENABLED
2154
2155	@classmethod
2156	@if_aiter_supported
2157	def is_fusion_moe_shared_experts_enabled(cls) -> bool:
2158	return cls.is_fused_moe_enabled() and cls._MOE_SHARED_EXPERTS_ENABLED
2159
2160	@classmethod
2161	@if_aiter_supported
2162	def is_fused_moe_situv2_enabled(cls) -> bool:
2163	"""True when a low-precision (a8w4/a4w4) SiTUv2 activation is selected."""
2164	return cls.is_fused_moe_enabled() and cls._MOE_SITUV2 != "a16w4"
2165
2166	@classmethod
2167	@if_aiter_supported
2168	def get_fused_moe_situv2_activation(cls) -> str:
2169	"""SiTUv2 activation dtype name: a16w4, a8w4 or a4w4."""
2170	return cls._MOE_SITUV2
2171
2172	@classmethod
2173	@if_aiter_supported
2174	def is_fused_moe_situv2_gate_up_interleaved(cls) -> bool:
2175	"""Whether the SiTUv2 kernels expect gate/up-interleaved w13.
2176
2177	The a8w4 FlyDSL stage-1 kernels (_gui_) are interleaved; a16w4 and
2178	a4w4 are separated. Weight shuffle (oracle/mxfp4.py) and gate_mode
2179	(experts/rocm_aiter_moe.py) must both use this.
2180	"""
2181	return cls._MOE_SITUV2 == "a8w4"
2182
2183	@classmethod
2184	@if_aiter_supported
2185	def topk_softmax_supports_fused_sigmoid(cls) -> bool:
2186	"""Check if topk_softmax supports fused shared expert activation."""
2187	if cls._TOPK_SOFTMAX_FUSED_SIGMOID is None:
2188	try:
2189	import inspect
2190
2191	from aiter import topk_softmax
2192
2193	params = inspect.signature(topk_softmax).parameters
2194	if "num_shared_experts" in params:
2195	cls._TOPK_SOFTMAX_FUSED_SIGMOID = True
2196	else:
2197	# @compile_ops wrapper loses the original signature.
2198	# Fall back to the torch custom op schema.
2199	import torch
2200
2201	schema = getattr(
2202	getattr(torch.ops.aiter, "topk_softmax", None), "default", None
2203	)
2204	schema_str = str(getattr(schema, "_schema", ""))
2205	cls._TOPK_SOFTMAX_FUSED_SIGMOID = "num_shared_experts" in schema_str
2206	except (ImportError, ValueError):
2207	cls._TOPK_SOFTMAX_FUSED_SIGMOID = False
2208	return cls._TOPK_SOFTMAX_FUSED_SIGMOID
2209
2210	@classmethod
2211	@if_aiter_supported
2212	def fuse_sigmoid_in_kernel(cls, aiter_topK_meta_data: object) -> bool:
2213	"""Whether fused shared-expert sigmoid in the topk kernel is usable.
2214
2215	Combines the cached static capability checks (FSE enabled, fused-moe
2216	enabled, topk_softmax supports fused sigmoid) with the runtime
2217	readiness check (topK meta-data buffer initialized).
2218
2219	``aiter_topK_meta_data`` is accepted as a parameter rather than
2220	imported internally so callers cannot hit initialization-order
2221	issues where the module-level global has not been set yet.
2222	"""
2223	return (
2224	cls.is_fusion_moe_shared_experts_enabled()
2225	and cls.topk_softmax_supports_fused_sigmoid()
2226	and aiter_topK_meta_data is not None
2227	)
2228
2229	@classmethod
2230	@if_aiter_supported
2231	def is_mla_enabled(cls) -> bool:
2232	return cls._AITER_ENABLED and cls._MLA_ENABLED
2233
2234	@classmethod
2235	@if_aiter_supported
2236	def mla_decode_supports_non_causal(cls) -> bool:
2237	"""Whether installed aiter.mla.mla_decode_fwd accepts `causal`.
2238
2239	Added in aiter v0.1.20. A missing argument is a hard error rather
2240	than a boolean the caller can use to pick another MLA backend.
2241	"""
2242	import inspect
2243
2244	from aiter.mla import mla_decode_fwd
2245
2246	if "causal" not in inspect.signature(mla_decode_fwd).parameters:
2247	raise RuntimeError(
2248	"ROCM_AITER_MLA requires aiter.mla.mla_decode_fwd(..., causal=). "
2249	"The installed aiter is causal-only; upgrade aiter rather than "
2250	"falling back to another MLA backend."
2251	)
2252	return True
2253
2254	@classmethod
2255	@if_aiter_supported
2256	def is_mha_enabled(cls) -> bool:
2257	return cls._AITER_ENABLED and cls._MHA_ENABLED
2258
2259	@classmethod
2260	@if_aiter_supported
2261	def is_custom_all_reduce_enabled(cls) -> bool:
2262	return cls._AITER_ENABLED and cls._CUSTOM_ALL_REDUCE_ENABLED
2263
2264	@classmethod
2265	@if_aiter_supported
2266	def is_shuffle_kv_cache_enabled(cls) -> bool:
2267	return cls._SHUFFLE_KV_CACHE_ENABLED
2268
2269	@classmethod
2270	@if_aiter_supported
2271	def fused_qknorm_idxrqknorm_enabled(
2272	cls,
2273	kv_cache_dtype: str,
2274	k_scale: torch.Tensor \| None = None,
2275	v_scale: torch.Tensor \| None = None,
2276	) -> bool:
2277	"""Whether MiniMax-M3 can use AITER's consolidated fused QK-norm.
2278
2279	Requires AITER to be installed and enabled. bf16 (``auto``) always
2280	qualifies; fp8 e4m3 also needs K/V scales. Other cache dtypes fall
2281	back to vLLM's fused kernel.
2282	"""
2283	if not cls._AITER_ENABLED:
2284	return False
2285	if kv_cache_dtype == "auto":
2286	return True
2287	if kv_cache_dtype in ("fp8", "fp8_e4m3"):
2288	return k_scale is not None and v_scale is not None
2289	return False
2290
2291	@classmethod
2292	def fused_qknorm_idxrqknorm(
2293	cls,
2294	qkv: torch.Tensor,
2295	q_norm_weight: torch.Tensor,
2296	k_norm_weight: torch.Tensor,
2297	cos_sin_cache: torch.Tensor,
2298	positions: torch.Tensor,
2299	num_heads: int,
2300	num_kv_heads: int,
2301	rotary_dim: int,
2302	eps: float,
2303	slot_mapping: torch.Tensor,
2304	kv_cache_k: torch.Tensor,
2305	kv_cache_v: torch.Tensor,
2306	q_out: torch.Tensor,
2307	kv_cache_dtype: str,
2308	k_scale: torch.Tensor \| None = None,
2309	v_scale: torch.Tensor \| None = None,
2310	index_q_norm_weight: torch.Tensor \| None = None,
2311	index_k_norm_weight: torch.Tensor \| None = None,
2312	num_index_heads: int = 0,
2313	index_cache: torch.Tensor \| None = None,
2314	index_q_out: torch.Tensor \| None = None,
2315	index_slot_mapping: torch.Tensor \| None = None,
2316	skip_index_branch: bool = False,
2317	) -> None:
2318	"""Run consolidated MiniMax-M3 QK-norm fusion.
2319
2320	Callers must check ``fused_qknorm_idxrqknorm_enabled`` first. Runtime
2321	failures from the AITER op are deliberately propagated.
2322	"""
2323	if kv_cache_dtype == "auto":
2324	aiter_kv_cache_dtype = "auto"
2325	aiter_k_scale = None
2326	aiter_v_scale = None
2327	elif kv_cache_dtype in ("fp8", "fp8_e4m3"):
2328	aiter_kv_cache_dtype = "fp8_e4m3_static"
2329	aiter_k_scale = k_scale
2330	aiter_v_scale = v_scale
2331	else:
2332	raise ValueError(
2333	"AITER fused QK-norm requires kv_cache_dtype 'auto', 'fp8', "
2334	f"or 'fp8_e4m3', got {kv_cache_dtype!r}"
2335	)
2336	aiter_index_cache_dtype = "fp8" if _is_fp8_e4m3_tensor(index_cache) else "auto"
2337
2338	from aiter import fused_qknorm_idxrqknorm
2339
2340	fused_qknorm_idxrqknorm(
2341	qkv,
2342	q_norm_weight,
2343	k_norm_weight,
2344	cos_sin_cache,
2345	positions,
2346	num_heads,
2347	num_kv_heads,
2348	rotary_dim,
2349	eps,
2350	index_q_norm_weight=index_q_norm_weight,
2351	index_k_norm_weight=index_k_norm_weight,
2352	num_index_heads=num_index_heads,
2353	slot_mapping=slot_mapping,
2354	kv_cache_k=kv_cache_k,
2355	kv_cache_v=kv_cache_v,
2356	index_cache=index_cache,
2357	block_size=16,
2358	q_out=q_out,
2359	index_q_out=index_q_out,
2360	index_slot_mapping=index_slot_mapping,
2361	kv_cache_dtype=aiter_kv_cache_dtype,
2362	index_cache_dtype=aiter_index_cache_dtype,
2363	k_scale=aiter_k_scale,
2364	v_scale=aiter_v_scale,
2365	asm_layout=True,
2366	skip_index_branch=skip_index_branch,
2367	)
2368
2369	@classmethod
2370	@if_aiter_supported
2371	def is_triton_unified_attn_enabled(cls) -> bool:
2372	return cls._AITER_ENABLED and cls._TRITON_UNIFIED_ATTN_ENABLED
2373
2374	@classmethod
2375	@if_aiter_supported
2376	def is_fp8bmm_enabled(cls) -> bool:
2377	return cls._AITER_ENABLED and cls._FP8BMM_ENABLED
2378
2379	@classmethod
2380	@if_aiter_supported
2381	def is_fp4bmm_enabled(cls) -> bool:
2382	from vllm.platforms.rocm import get_cdna_version
2383
2384	# TODO GFX1250: Enable for cdna 4+ when aiter supports batched_gemm_a16wfp4 on gfx1250
2385	return cls._AITER_ENABLED and cls._FP4BMM_ENABLED and get_cdna_version() == 4
2386
2387	@classmethod
2388	@if_aiter_supported
2389	def is_linear_hipbmm_enabled(cls) -> bool:
2390	from vllm.platforms.rocm import get_cdna_version
2391
2392	return (
2393	cls.is_linear_enabled()
2394	and (get_cdna_version() > 2)
2395	and cls._LINEAR_HIPBMM_ENABLED
2396	)
2397
2398	@classmethod
2399	@if_aiter_supported
2400	def is_asm_fp4_gemm_dynamic_quant_enabled(cls) -> bool:
2401	from vllm.platforms.rocm import on_gfx950
2402
2403	return cls._AITER_ENABLED and cls._FP4_GEMM_DYNAMIC_QUANT_ASM and on_gfx950()
2404
2405	@classmethod
2406	@if_aiter_supported
2407	def is_triton_rotary_embed_enabled(cls) -> bool:
2408	return cls._AITER_ENABLED and cls._TRITON_ROTARY_EMBED
2409
2410	@classmethod
2411	@if_aiter_supported
2412	def is_triton_sparse_mla_enabled(cls) -> bool:
2413	if not cls._TRITON_SPARSE_MLA:
2414	return False
2415	from vllm.platforms.rocm import on_gfx950
2416
2417	if not cls._AITER_ENABLED:
2418	reason = "VLLM_ROCM_USE_AITER is off"
2419	elif not on_gfx950():
2420	reason = "the kernel is gfx950-only"
2421	else:
2422	return True
2423	logger.warning_once(
2424	"VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA is set, but %s; ignoring it.",
2425	reason,
2426	)
2427	return False
2428
2429	@classmethod
2430	@if_aiter_supported
2431	def is_triton_gemm_enabled(cls) -> bool:
2432	return cls._AITER_ENABLED and cls._TRITON_UNQUANT_GEMM
2433
2434	@classmethod
2435	@if_aiter_supported
2436	def is_tgemm_enabled(cls) -> bool:
2437	from vllm.platforms.rocm import on_gfx950
2438
2439	return cls.is_linear_enabled() and on_gfx950()
2440
2441	@classmethod
2442	def get_aiter_allreduce(cls):
2443	"""Return the TP device communicator's AITER custom-allreduce if it has
2444	one, return None otherwise
2445	"""
2446	from vllm.distributed.device_communicators.aiter_custom_all_reduce import (
2447	AiterCustomAllreduce,
2448	)
2449	from vllm.distributed.parallel_state import get_tp_group
2450
2451	device_comm = get_tp_group().device_communicator
2452	aiter_ar_comm = getattr(device_comm, "aiter_ar_comm", None)
2453	return (
2454	aiter_ar_comm if isinstance(aiter_ar_comm, AiterCustomAllreduce) else None
2455	)
2456
2457	@staticmethod
2458	def _gdn_triton_kernels_importable() -> bool:
2459	try:
2460	import aiter.ops.triton.causal_conv1d_update_single_token # noqa: F401
2461	import aiter.ops.triton.gated_delta_net # noqa: F401
2462	from aiter.ops.triton.quant import ( # noqa: F401
2463	fused_rms_gated_fp8_group_quant,
2464	)
2465
2466	return True
2467	except (ImportError, ModuleNotFoundError):
2468	return False
2469
2470	@classmethod
2471	@if_aiter_supported
2472	def are_gdn_triton_kernels_available(cls) -> bool:
2473	"""Check if AITER Triton kernels for GDN attention are importable.
2474
2475	These are optional Triton kernels (conv1d fast-path, gated delta net)
2476	used by GatedDeltaNetAttention's decode fast-path. They may be absent
2477	in older aiter builds.
2478	"""
2479	return cls._AITER_ENABLED and cls._gdn_triton_kernels_importable()
2480
2481	@classmethod
2482	def is_rdna_gdn_triton_kernels_available(cls) -> bool:
2483	"""RDNA4 (gfx12) analog of are_gdn_triton_kernels_available()."""
2484	return cls.is_rdna_aiter_enabled() and cls._gdn_triton_kernels_importable()
2485
2486	@classmethod
2487	@if_aiter_supported
2488	@functools.cache
2489	def fused_moe_supports_gate_mode(cls) -> bool:
2490	"""Probe whether the installed aiter.fused_moe accepts `gate_mode`.
2491
2492	Added in https://github.com/ROCm/aiter/pull/3123 (>=0.1.14).
2493	Builds with older AITER must omit this argument.
2494	"""
2495	import inspect
2496
2497	from aiter.fused_moe import fused_moe
2498
2499	return "gate_mode" in inspect.signature(fused_moe).parameters
2500
2501	@classmethod
2502	@if_aiter_supported
2503	@functools.cache
2504	def fused_moe_supports_quant_dtype_a(cls) -> bool:
2505	"""Probe whether the installed aiter.fused_moe accepts `quant_dtype_a`.
2506
2507	Added in https://github.com/ROCm/aiter/pull/5439 (unreleased at
2508	merge time). Older AITER can't override the activation quant dtype.
2509	"""
2510	import inspect
2511
2512	from aiter.fused_moe import fused_moe
2513
2514	return "quant_dtype_a" in inspect.signature(fused_moe).parameters
2515
2516	@staticmethod
2517	def _probe_dsv4_i384_fhmoe_capability(num_tokens: int) -> bool:
2518	"""Probe AITER's CSV-backed DSV4 native-I384 FHMoE contract."""
2519	if type(num_tokens) is not int or num_tokens <= 0:
2520	return False
2521
2522	try:
2523	import inspect
2524
2525	import aiter.fhmoe as fhmoe
2526	from aiter.fused_moe import fused_moe
2527
2528	params = inspect.signature(fused_moe).parameters
2529	supports_dsv4_i384_fhmoe = getattr(fhmoe, "supports_dsv4_i384_fhmoe", None)
2530	if not callable(supports_dsv4_i384_fhmoe):
2531	return False
2532	supports_num_tokens = supports_dsv4_i384_fhmoe(num_tokens)
2533	except Exception:
2534	return False
2535
2536	return supports_num_tokens is True and all(
2537	name in params
2538	for name in (
2539	"shared_w1",
2540	"shared_w2",
2541	"shared_w1_scale",
2542	"shared_w2_scale",
2543	"shared_expert_id",
2544	)
2545	)
2546
2547	@classmethod
2548	@if_aiter_supported
2549	@functools.cache
2550	def fused_moe_supports_heterogeneous_shared_expert(cls, num_tokens: int) -> bool:
2551	"""Whether AITER has DSV4 native-I384 configs through the given M."""
2552	return cls._probe_dsv4_i384_fhmoe_capability(num_tokens)
2553
2554	@classmethod
2555	@if_aiter_supported
2556	def is_indexer_top_k_enabled(cls) -> bool:
2557	"""gfx950 is the only arch with tuned AITER indexer top-k kernels."""
2558	from vllm.platforms.rocm import on_gfx950
2559
2560	return cls._AITER_ENABLED and on_gfx950()
2561
2562	@classmethod
2563	def is_indexer_top_k_supported(
2564	cls,
2565	*,
2566	is_prefill: bool,
2567	compress_ratio: int,
2568	num_rows: int,
2569	max_valid_seq_len: int \| None = None,
2570	) -> bool:
2571	"""Whether AITER's sparse indexer top-k beats the in-tree kernel for
2572	this shape."""
2573	if compress_ratio <= 1 or not cls.is_indexer_top_k_enabled():
2574	return False
2575
2576	if not is_prefill:
2577	assert max_valid_seq_len is not None
2578	# AITER v0.1.19 decode is one-block only. This measured gfx950
2579	# FP32/k=1024 compressed-row boundary is independent of the native
2580	# split-count boundary in sampler.cu.
2581	if (
2582	num_rows <= _GFX950_C4A_NATIVE_MAX_ROWS
2583	and max_valid_seq_len > _GFX950_C4A_AITER_MAX_COMPRESSED_SEQ_LEN
2584	):
2585	return False
2586
2587	return True
2588
2589	@staticmethod
2590	def dsv4_indexer_prefers_native_top_k(
2591	*,
2592	num_rows: int,
2593	num_columns: int,
2594	topk_tokens: int,
2595	) -> bool:
2596	"""The in-tree decode kernel's measured advantage window over AITER.
2597	Tuned on DSV4's compressed-KV logits, so it applies to that indexer
2598	only."""
2599	return (
2600	topk_tokens == 512
2601	and 0 < num_rows <= 384
2602	and num_columns <= _GFX950_DSV4_NATIVE_MAX_COLUMNS
2603	)
2604
2605	@staticmethod
2606	def indexer_top_k_decode(
2607	logits: torch.Tensor,
2608	next_n: int,
2609	seq_lens: torch.Tensor,
2610	topk_indices: torch.Tensor,
2611	topk_tokens: int,
2612	) -> None:
2613	from aiter.ops.topk import top_k_per_row_decode
2614
2615	top_k_per_row_decode(
2616	logits,
2617	next_n,
2618	seq_lens,
2619	topk_indices,
2620	logits.shape[0],
2621	logits.stride(0),
2622	logits.stride(1),
2623	k=topk_tokens,
2624	)
2625
2626	@staticmethod
2627	def indexer_top_k_prefill(
2628	logits: torch.Tensor,
2629	row_starts: torch.Tensor,
2630	row_ends: torch.Tensor,
2631	indices: torch.Tensor,
2632	topk_tokens: int,
2633	) -> None:
2634	from aiter.ops.topk import top_k_per_row_prefill
2635
2636	top_k_per_row_prefill(
2637	logits,
2638	row_starts,
2639	row_ends,
2640	indices,
2641	None,
2642	logits.shape[0],
2643	logits.stride(0),
2644	logits.stride(1),
2645	k=topk_tokens,
2646	)
2647
2648	@staticmethod
2649	def register_ops_once() -> None:
2650	global _OPS_REGISTERED
2651
2652	if not (
2653	is_aiter_found_and_supported() or is_aiter_found_and_supported_on_rdna4()
2654	):
2655	if not current_platform.is_rocm():
2656	return
2657
2658	from vllm.platforms.rocm import on_gfx11, on_gfx950
2659
2660	# This op has self-contained Triton/C++ implementations on gfx11
2661	# and gfx950. Only its optional top-k fast path comes from aiter.
2662	if (on_gfx11() or on_gfx950()) and not _OPS_REGISTERED:
2663	direct_register_custom_op(
2664	op_name="rocm_aiter_sparse_attn_indexer",
2665	op_func=rocm_aiter_sparse_attn_indexer,
2666	mutates_args=["topk_indices_buffer", "candidate_blocks"],
2667	fake_impl=rocm_aiter_sparse_attn_indexer_fake,
2668	dispatch_key=current_platform.dispatch_key,
2669	)
2670	_OPS_REGISTERED = True
2671	return
2672
2673	if not _OPS_REGISTERED:
2674	# register all the custom ops here
2675	direct_register_custom_op(
2676	op_name="rocm_aiter_asm_moe_tkw1",
2677	op_func=_rocm_aiter_asm_moe_tkw1_impl,
2678	mutates_args=[],
2679	fake_impl=_rocm_aiter_asm_moe_tkw1_fake,
2680	dispatch_key=current_platform.dispatch_key,
2681	)
2682
2683	direct_register_custom_op(
2684	op_name="rocm_aiter_fused_moe",
2685	op_func=_rocm_aiter_fused_moe_impl,
2686	mutates_args=[],
2687	fake_impl=_rocm_aiter_fused_moe_fake,
2688	dispatch_key=current_platform.dispatch_key,
2689	)
2690
2691	direct_register_custom_op(
2692	op_name="rocm_aiter_topk_softmax",
2693	op_func=_rocm_aiter_topk_softmax_impl,
2694	mutates_args=["topk_weights", "topk_indices", "token_expert_indices"],
2695	dispatch_key=current_platform.dispatch_key,
2696	)
2697
2698	direct_register_custom_op(
2699	op_name="rocm_aiter_topk_sigmoid",
2700	op_func=_rocm_aiter_topk_sigmoid_impl,
2701	mutates_args=["topk_weights", "topk_indices"],
2702	dispatch_key=current_platform.dispatch_key,
2703	)
2704
2705	direct_register_custom_op(
2706	op_name="rocm_aiter_biased_grouped_topk",
2707	op_func=_rocm_aiter_biased_grouped_topk_impl,
2708	mutates_args=["topk_weights", "topk_ids"],
2709	dispatch_key=current_platform.dispatch_key,
2710	)
2711
2712	direct_register_custom_op(
2713	op_name="rocm_aiter_grouped_topk",
2714	op_func=_rocm_aiter_grouped_topk_impl,
2715	mutates_args=["topk_weights", "topk_ids"],
2716	dispatch_key=current_platform.dispatch_key,
2717	)
2718
2719	direct_register_custom_op(
2720	op_name="rocm_aiter_fused_topk",
2721	op_func=_rocm_aiter_fused_topk_impl,
2722	mutates_args=[],
2723	fake_impl=_rocm_aiter_fused_topk_fake,
2724	dispatch_key=current_platform.dispatch_key,
2725	)
2726
2727	direct_register_custom_op(
2728	op_name="rocm_aiter_mla_decode_fwd",
2729	op_func=_rocm_aiter_mla_decode_fwd_impl,
2730	mutates_args=["o"],
2731	)
2732
2733	direct_register_custom_op(
2734	op_name="rocm_aiter_mla_decode_fwd_lse",
2735	op_func=_rocm_aiter_mla_decode_fwd_lse_impl,
2736	mutates_args=["o"],
2737	fake_impl=_rocm_aiter_mla_decode_fwd_lse_fake,
2738	)
2739
2740	direct_register_custom_op(
2741	op_name="rocm_aiter_w8a8_gemm",
2742	op_func=_rocm_aiter_w8a8_gemm_impl,
2743	fake_impl=_rocm_aiter_w8a8_gemm_fake,
2744	)
2745
2746	direct_register_custom_op(
2747	op_name="_rocm_aiter_preshuffled_per_token_w8a8_gemm",
2748	op_func=_rocm_aiter_preshuffled_per_token_w8a8_gemm_impl,
2749	fake_impl=_rocm_aiter_preshuffled_per_token_w8a8_gemm_fake,
2750	)
2751
2752	direct_register_custom_op(
2753	op_name="rocm_aiter_hipb_mm_fp8",
2754	op_func=_rocm_aiter_hipb_mm_fp8_impl,
2755	fake_impl=_rocm_aiter_hipb_mm_fp8_fake,
2756	)
2757
2758	direct_register_custom_op(
2759	op_name="rocm_aiter_triton_gemm_a8w8_blockscale",
2760	op_func=_rocm_aiter_triton_gemm_a8w8_blockscale_impl,
2761	fake_impl=_rocm_aiter_triton_gemm_a8w8_blockscale_fake,
2762	)
2763
2764	direct_register_custom_op(
2765	op_name="rocm_aiter_gemm_a8w8_blockscale",
2766	op_func=_rocm_aiter_gemm_a8w8_blockscale_impl,
2767	fake_impl=_rocm_aiter_gemm_a8w8_blockscale_fake,
2768	)
2769
2770	direct_register_custom_op(
2771	op_name="rocm_aiter_gemm_a8w8_blockscale_bpreshuffle",
2772	op_func=_rocm_aiter_gemm_a8w8_blockscale_bpreshuffle_impl,
2773	fake_impl=_rocm_aiter_gemm_a8w8_blockscale_bpreshuffle_fake,
2774	)
2775
2776	direct_register_custom_op(
2777	op_name="rocm_aiter_rmsnorm_fused_dynamic_quant",
2778	op_func=_rocm_aiter_rmsnorm_fused_dynamic_quant_impl,
2779	fake_impl=_rocm_aiter_rmsnorm_fused_dynamic_quant_fake,
2780	dispatch_key=current_platform.dispatch_key,
2781	)
2782
2783	direct_register_custom_op(
2784	op_name="rocm_aiter_rmsnorm_fused_add_dynamic_quant",
2785	op_func=_rocm_aiter_rmsnorm_fused_add_dynamic_quant_impl,
2786	fake_impl=_rocm_aiter_rmsnorm_fused_add_dynamic_quant_fake,
2787	dispatch_key=current_platform.dispatch_key,
2788	)
2789
2790	direct_register_custom_op(
2791	op_name="rocm_aiter_rmsnorm_fp8_group_quant",
2792	op_func=_rocm_aiter_rmsnorm_fp8_group_quant_impl,
2793	fake_impl=_rocm_aiter_rmsnorm_fp8_group_quant_fake,
2794	)
2795
2796	direct_register_custom_op(
2797	op_name="rocm_aiter_fused_rms_gated_fp8_group_quant",
2798	op_func=_rocm_aiter_fused_rms_gated_fp8_group_quant_impl,
2799	fake_impl=_rocm_aiter_fused_rms_gated_fp8_group_quant_fake,
2800	)
2801
2802	direct_register_custom_op(
2803	op_name="rocm_aiter_rmsnorm_with_add_fp8_group_quant",
2804	op_func=_rocm_aiter_rmsnorm_with_add_fp8_group_quant_impl,
2805	fake_impl=_rocm_aiter_rmsnorm_with_add_fp8_group_quant_fake,
2806	)
2807
2808	direct_register_custom_op(
2809	op_name="rocm_aiter_act_mul_and_fp8_group_quant",
2810	op_func=_rocm_aiter_act_mul_and_fp8_group_quant_impl,
2811	fake_impl=_rocm_aiter_act_mul_and_fp8_group_quant_fake,
2812	)
2813
2814	direct_register_custom_op(
2815	op_name="rocm_aiter_triton_add_rmsnorm_pad",
2816	op_func=_rocm_aiter_triton_add_rmsnorm_pad_impl,
2817	fake_impl=_rocm_aiter_triton_add_rmsnorm_pad_fake,
2818	dispatch_key=current_platform.dispatch_key,
2819	)
2820
2821	direct_register_custom_op(
2822	op_name="rocm_aiter_group_fp8_quant",
2823	op_func=_rocm_aiter_group_fp8_quant_impl,
2824	fake_impl=_rocm_aiter_group_fp8_quant_fake,
2825	)
2826
2827	direct_register_custom_op(
2828	op_name="rocm_aiter_per_tensor_quant",
2829	op_func=_rocm_aiter_per_tensor_quant_impl,
2830	mutates_args=["out", "scale"],
2831	dispatch_key=current_platform.dispatch_key,
2832	)
2833
2834	direct_register_custom_op(
2835	op_name="rocm_aiter_per_token_quant",
2836	op_func=_rocm_aiter_per_token_quant_impl,
2837	fake_impl=_rocm_aiter_per_token_quant_fake,
2838	dispatch_key=current_platform.dispatch_key,
2839	)
2840
2841	direct_register_custom_op(
2842	op_name="rocm_aiter_sparse_attn_indexer",
2843	op_func=rocm_aiter_sparse_attn_indexer,
2844	mutates_args=["topk_indices_buffer", "candidate_blocks"],
2845	fake_impl=rocm_aiter_sparse_attn_indexer_fake,
2846	dispatch_key=current_platform.dispatch_key,
2847	)
2848
2849	direct_register_custom_op(
2850	op_name="aiter_fp8_attn_wrapper",
2851	op_func=_rocm_aiter_fp8_attn_impl,
2852	mutates_args=[],
2853	fake_impl=_rocm_aiter_fp8_attn_fake,
2854	dispatch_key=current_platform.dispatch_key,
2855	)
2856
2857	direct_register_custom_op(
2858	op_name="rocm_aiter_gemm_a8wfp4",
2859	op_func=_rocm_aiter_gemm_a8wfp4_impl,
2860	mutates_args=[],
2861	fake_impl=_rocm_aiter_gemm_a8wfp4_fake,
2862	dispatch_key=current_platform.dispatch_key,
2863	)
2864
2865	# Register rocm aiter rotary embedding custom op
2866	direct_register_custom_op(
2867	op_name="rocm_aiter_triton_rotary_embedding",
2868	op_func=_triton_rotary_embedding_impl,
2869	mutates_args=["query", "key"], # These tensors are modified in-place
2870	)
2871
2872	direct_register_custom_op(
2873	op_name="rocm_aiter_fused_allreduce_rmsnorm",
2874	op_func=_rocm_aiter_fused_allreduce_rmsnorm_impl,
2875	fake_impl=_rocm_aiter_fused_allreduce_rmsnorm_fake,
2876	)
2877
2878	direct_register_custom_op(
2879	op_name="rocm_aiter_fused_allreduce_rmsnorm_quant_per_group",
2880	op_func=(_rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_impl),
2881	fake_impl=(_rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_fake),
2882	)
2883
2884	direct_register_custom_op(
2885	op_name="rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm", # noqa: E501
2886	op_func=_rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_impl, # noqa: E501
2887	fake_impl=_rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_fake, # noqa: E501
2888	)
2889
2890	direct_register_custom_op(
2891	op_name="fused_mla_dual_rms_norm",
2892	op_func=_fused_mla_dual_rms_norm_impl,
2893	mutates_args=[],
2894	fake_impl=_fused_mla_dual_rms_norm_fake,
2895	)
2896
2897	direct_register_custom_op(
2898	op_name="fused_mla_dual_rms_norm_per_token_quant",
2899	op_func=_fused_mla_dual_rms_norm_per_token_quant_impl,
2900	mutates_args=[],
2901	fake_impl=_fused_mla_dual_rms_norm_per_token_quant_fake,
2902	)
2903
2904	direct_register_custom_op(
2905	op_name="fused_mla_dual_rms_norm_group_quant",
2906	op_func=_fused_mla_dual_rms_norm_group_quant_impl,
2907	mutates_args=[],
2908	fake_impl=_fused_mla_dual_rms_norm_group_quant_fake,
2909	)
2910
2911	_OPS_REGISTERED = True
2912
2913	@staticmethod
2914	def get_rmsnorm_fused_add_dynamic_quant_op() -> OpOverload:
2915	return torch.ops.vllm.rocm_aiter_rmsnorm_fused_add_dynamic_quant.default
2916
2917	@staticmethod
2918	def get_rmsnorm_fused_dynamic_quant_op() -> OpOverload:
2919	return torch.ops.vllm.rocm_aiter_rmsnorm_fused_dynamic_quant.default
2920
2921	@staticmethod
2922	def get_rmsnorm_group_fused_quant_op() -> OpOverload:
2923	return torch.ops.vllm.rocm_aiter_rmsnorm_fp8_group_quant.default
2924
2925	@staticmethod
2926	def get_fused_rms_gated_fp8_group_quant_op() -> OpOverload:
2927	"""Return the fused gated-RMSNorm + FP8 group quant custom op."""
2928	return torch.ops.vllm.rocm_aiter_fused_rms_gated_fp8_group_quant.default
2929
2930	@staticmethod
2931	def get_rmsnorm_group_add_fused_quant_op() -> OpOverload:
2932	return torch.ops.vllm.rocm_aiter_rmsnorm_with_add_fp8_group_quant.default
2933
2934	@staticmethod
2935	def get_per_token_quant_op() -> OpOverload:
2936	return torch.ops.vllm.rocm_aiter_per_token_quant.default
2937
2938	@staticmethod
2939	def get_group_quant_op() -> OpOverload:
2940	return torch.ops.vllm.rocm_aiter_group_fp8_quant.default
2941
2942	@staticmethod
2943	def get_act_mul_fused_fp8_group_quant_op() -> OpOverload:
2944	return torch.ops.vllm.rocm_aiter_act_mul_and_fp8_group_quant.default
2945
2946	@staticmethod
2947	def get_triton_add_rmsnorm_pad_op() -> OpOverload:
2948	return torch.ops.vllm.rocm_aiter_triton_add_rmsnorm_pad.default
2949
2950	@staticmethod
2951	def get_triton_rotary_embedding_op() -> OpOverload:
2952	return torch.ops.vllm.rocm_aiter_triton_rotary_embedding.default
2953
2954	@staticmethod
2955	def get_fused_allreduce_rmsnorm_op() -> OpOverload:
2956	return torch.ops.vllm.rocm_aiter_fused_allreduce_rmsnorm.default
2957
2958	@staticmethod
2959	def get_fused_allreduce_rmsnorm_quant_per_group_op() -> OpOverload:
2960	return torch.ops.vllm.rocm_aiter_fused_allreduce_rmsnorm_quant_per_group.default
2961
2962	@staticmethod
2963	def get_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_op() -> OpOverload: # noqa: E501
2964	return torch.ops.vllm.rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm.default # noqa: E501
2965
2966	@staticmethod
2967	def get_fused_mla_dual_rms_norm_op() -> OpOverload:
2968	return torch.ops.vllm.fused_mla_dual_rms_norm.default
2969
2970	@staticmethod
2971	def get_fused_mla_dual_rms_norm_per_token_quant_op() -> OpOverload:
2972	return torch.ops.vllm.fused_mla_dual_rms_norm_per_token_quant.default
2973
2974	@staticmethod
2975	def fused_qk_rmsnorm_group_quant(
2976	q: torch.Tensor,
2977	q_weight: torch.Tensor,
2978	q_epsilon: float,
2979	kv: torch.Tensor,
2980	kv_weight: torch.Tensor,
2981	kv_epsilon: float,
2982	group_size: int = 128,
2983	transpose_scale: bool = False,
2984	) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
2985	"""Fused q/kv RMSNorm + per-1x128 FP8 group quant on q via AITER."""
2986	return torch.ops.vllm.fused_mla_dual_rms_norm_group_quant(
2987	q,
2988	q_weight,
2989	kv,
2990	kv_weight,
2991	q_epsilon,
2992	kv_epsilon,
2993	group_size,
2994	transpose_scale,
2995	)
2996
2997	@staticmethod
2998	def rms_norm(
2999	x: torch.Tensor,
3000	weight: torch.Tensor,
3001	epsilon: float,
3002	) -> torch.Tensor:
3003	"""RMSNorm via AITER kernel."""
3004	import aiter
3005
3006	return aiter.rms_norm(x, weight, epsilon)
3007
3008	@staticmethod
3009	def rms_norm2d_with_add(
3010	x: torch.Tensor,
3011	residual: torch.Tensor,
3012	weight: torch.Tensor,
3013	epsilon: float,
3014	) -> tuple[torch.Tensor, torch.Tensor]:
3015	"""Fused residual-add + RMSNorm via AITER kernel.
3016
3017	Returns (normalized_output, residual_sum).
3018	"""
3019	import aiter
3020
3021	out = torch.empty_like(x)
3022	residual_out = torch.empty_like(x)
3023	aiter.rmsnorm2d_fwd_with_add(out, x, residual, residual_out, weight, epsilon, 0)
3024	return out, residual_out
3025
3026	@staticmethod
3027	def w8a8_gemm(
3028	A: torch.Tensor,
3029	B: torch.Tensor,
3030	As: torch.Tensor,
3031	Bs: torch.Tensor,
3032	bias: torch.Tensor \| None = None,
3033	output_dtype: torch.dtype = torch.float16,
3034	) -> torch.Tensor:
3035	return torch.ops.vllm.rocm_aiter_w8a8_gemm(A, B, As, Bs, bias, output_dtype)
3036
3037	@staticmethod
3038	def preshuffled_per_token_w8a8_gemm(
3039	A: torch.Tensor,
3040	B: torch.Tensor,
3041	As: torch.Tensor,
3042	Bs: torch.Tensor,
3043	bias: torch.Tensor \| None = None,
3044	output_dtype: torch.dtype = torch.float16,
3045	) -> torch.Tensor:
3046	return torch.ops.vllm._rocm_aiter_preshuffled_per_token_w8a8_gemm(
3047	A, B, As, Bs, bias, output_dtype
3048	)
3049
3050	@staticmethod
3051	def hipb_mm_fp8(
3052	A: torch.Tensor,
3053	B: torch.Tensor,
3054	As: torch.Tensor,
3055	Bs: torch.Tensor,
3056	bias: torch.Tensor \| None = None,
3057	output_dtype: torch.dtype = torch.bfloat16,
3058	) -> torch.Tensor:
3059	return torch.ops.vllm.rocm_aiter_hipb_mm_fp8(A, B, As, Bs, bias, output_dtype)
3060
3061	@staticmethod
3062	def triton_gemm_a8w8_blockscale(
3063	A: torch.Tensor,
3064	B: torch.Tensor,
3065	As: torch.Tensor,
3066	Bs: torch.Tensor,
3067	block_size: list[int],
3068	output_dtype: torch.dtype = torch.float16,
3069	) -> torch.Tensor:
3070	return torch.ops.vllm.rocm_aiter_triton_gemm_a8w8_blockscale(
3071	A, B, As, Bs, output_dtype
3072	)
3073
3074	@staticmethod
3075	def gemm_a8w8_blockscale(
3076	A: torch.Tensor,
3077	B: torch.Tensor,
3078	As: torch.Tensor,
3079	Bs: torch.Tensor,
3080	block_size: list[int],
3081	output_dtype: torch.dtype = torch.float16,
3082	) -> torch.Tensor:
3083	return torch.ops.vllm.rocm_aiter_gemm_a8w8_blockscale(
3084	A, B, As, Bs, output_dtype
3085	)
3086
3087	@staticmethod
3088	def gemm_a8w8_blockscale_bpreshuffle(
3089	A: torch.Tensor,
3090	B: torch.Tensor,
3091	As: torch.Tensor,
3092	Bs: torch.Tensor,
3093	output_dtype: torch.dtype = torch.float16,
3094	) -> torch.Tensor:
3095	return torch.ops.vllm.rocm_aiter_gemm_a8w8_blockscale_bpreshuffle(
3096	A, B, As, Bs, output_dtype
3097	)
3098
3099	@staticmethod
3100	def fused_moe(
3101	hidden_states: torch.Tensor,
3102	w1: torch.Tensor,
3103	w2: torch.Tensor,
3104	topk_weight: torch.Tensor,
3105	topk_ids: torch.Tensor,
3106	expert_mask: torch.Tensor \| None = None,
3107	activation_method: int = 0,
3108	quant_method: int = 0,
3109	doweight_stage1: bool = False,
3110	w1_scale: torch.Tensor \| None = None,
3111	w2_scale: torch.Tensor \| None = None,
3112	a1_scale: torch.Tensor \| None = None,
3113	a2_scale: torch.Tensor \| None = None,
3114	num_local_tokens: torch.Tensor \| None = None,
3115	output_dtype: torch.dtype \| None = None,
3116	hidden_pad: int = 0,
3117	intermediate_pad: int = 0,
3118	gate_mode: str = "",
3119	bias1: torch.Tensor \| None = None,
3120	bias2: torch.Tensor \| None = None,
3121	moe_sorting_dispatch_policy: int = 0,
3122	swiglu_limit: float = 0.0,
3123	beta: float \| None = None,
3124	linear_beta: float \| None = None,
3125	shared_w1: torch.Tensor \| None = None,
3126	shared_w2: torch.Tensor \| None = None,
3127	shared_w1_scale: torch.Tensor \| None = None,
3128	shared_w2_scale: torch.Tensor \| None = None,
3129	shared_expert_id: int = -1,
3130	q_dtype_a: torch.dtype \| None = None,
3131	) -> torch.Tensor:
3132	return torch.ops.vllm.rocm_aiter_fused_moe(
3133	hidden_states,
3134	w1,
3135	w2,
3136	topk_weight,
3137	topk_ids,
3138	expert_mask,
3139	activation_method,
3140	quant_method,
3141	doweight_stage1,
3142	w1_scale,
3143	w2_scale,
3144	a1_scale,
3145	a2_scale,
3146	num_local_tokens,
3147	output_dtype,
3148	hidden_pad,
3149	intermediate_pad,
3150	gate_mode,
3151	bias1,
3152	bias2,
3153	moe_sorting_dispatch_policy,
3154	swiglu_limit,
3155	beta,
3156	linear_beta,
3157	shared_w1,
3158	shared_w2,
3159	shared_w1_scale,
3160	shared_w2_scale,
3161	shared_expert_id,
3162	q_dtype_a,
3163	)
3164
3165	@staticmethod
3166	def asm_moe_tkw1(
3167	hidden_states: torch.Tensor,
3168	w1: torch.Tensor,
3169	w2: torch.Tensor,
3170	topk_weights: torch.Tensor,
3171	topk_ids: torch.Tensor,
3172	fc1_scale: torch.Tensor \| None = None,
3173	fc2_scale: torch.Tensor \| None = None,
3174	fc1_smooth_scale: torch.Tensor \| None = None,
3175	fc2_smooth_scale: torch.Tensor \| None = None,
3176	a16: bool = False,
3177	per_tensor_quant_scale: torch.Tensor \| None = None,
3178	expert_mask: torch.Tensor \| None = None,
3179	activation_method: int = 0,
3180	) -> torch.Tensor:
3181	return torch.ops.vllm.rocm_aiter_asm_moe_tkw1(
3182	hidden_states,
3183	w1,
3184	w2,
3185	topk_weights,
3186	topk_ids,
3187	fc1_scale,
3188	fc2_scale,
3189	fc1_smooth_scale,
3190	fc2_smooth_scale,
3191	a16,
3192	per_tensor_quant_scale,
3193	expert_mask,
3194	activation_method,
3195	)
3196
3197	@staticmethod
3198	def topk_softmax(
3199	topk_weights: torch.Tensor,
3200	topk_indices: torch.Tensor,
3201	token_expert_indices: torch.Tensor,
3202	gating_output: torch.Tensor,
3203	renormalize: bool,
3204	num_shared_experts: int = 0,
3205	shared_expert_scoring_func: str = "",
3206	) -> tuple[torch.Tensor, ...]:
3207	torch.ops.vllm.rocm_aiter_topk_softmax(
3208	topk_weights,
3209	topk_indices,
3210	token_expert_indices,
3211	gating_output,
3212	renormalize,
3213	num_shared_experts,
3214	shared_expert_scoring_func,
3215	)
3216	return topk_weights, topk_indices
3217
3218	@staticmethod
3219	def topk_sigmoid(
3220	topk_weights: torch.Tensor,
3221	topk_indices: torch.Tensor,
3222	token_expert_indices: torch.Tensor,
3223	gating_output: torch.Tensor,
3224	renormalize: bool,
3225	) -> tuple[torch.Tensor, ...]:
3226	torch.ops.vllm.rocm_aiter_topk_sigmoid(
3227	topk_weights, topk_indices, gating_output
3228	)
3229	return topk_weights, topk_indices
3230
3231	@staticmethod
3232	def biased_grouped_topk(
3233	gating_output: torch.Tensor,
3234	correction_bias: torch.Tensor,
3235	topk_weights: torch.Tensor,
3236	topk_ids: torch.Tensor,
3237	num_expert_group: int,
3238	topk_group: int,
3239	need_renorm: bool,
3240	routed_scaling_factor: float = 1.0,
3241	) -> None:
3242	if correction_bias.dtype != gating_output.dtype:
3243	correction_bias = correction_bias.to(gating_output.dtype)
3244	torch.ops.vllm.rocm_aiter_biased_grouped_topk(
3245	gating_output,
3246	correction_bias,
3247	topk_weights,
3248	topk_ids,
3249	num_expert_group,
3250	topk_group,
3251	need_renorm,
3252	routed_scaling_factor,
3253	)
3254
3255	@staticmethod
3256	def grouped_topk(
3257	gating_output: torch.Tensor,
3258	topk_weights: torch.Tensor,
3259	topk_ids: torch.Tensor,
3260	num_expert_group: int,
3261	topk_group: int,
3262	need_renorm: bool,
3263	scoring_func: str = "softmax",
3264	routed_scaling_factor: float = 1.0,
3265	) -> None:
3266	torch.ops.vllm.rocm_aiter_grouped_topk(
3267	gating_output,
3268	topk_weights,
3269	topk_ids,
3270	num_expert_group,
3271	topk_group,
3272	need_renorm,
3273	scoring_func,
3274	routed_scaling_factor,
3275	)
3276
3277	@staticmethod
3278	def fused_topk(
3279	x: torch.Tensor,
3280	router_logits: torch.Tensor,
3281	top_k: int,
3282	gate_up: bool,
3283	) -> tuple[torch.Tensor, torch.Tensor]:
3284	return torch.ops.vllm.rocm_aiter_fused_topk(x, router_logits, top_k, gate_up)
3285
3286	@staticmethod
3287	def triton_sparse_mla_fwd(
3288	q: torch.Tensor,
3289	kv_buffer: torch.Tensor,
3290	o: torch.Tensor,
3291	sm_scale: float,
3292	kv_indptr: torch.Tensor,
3293	kv_indices: torch.Tensor,
3294	kv_lora_rank: int = 512,
3295	qk_rope_head_dim: int = 64,
3296	q_scale: torch.Tensor \| None = None,
3297	kv_scale: torch.Tensor \| None = None,
3298	attn_sink: torch.Tensor \| None = None,
3299	extra_kv_buffer: torch.Tensor \| None = None,
3300	extra_kv_indptr: torch.Tensor \| None = None,
3301	extra_kv_indices: torch.Tensor \| None = None,
3302	has_invalid: bool = True,
3303	) -> None:
3304	"""Sparse MLA read straight from the KV cache, for prefill and decode.
3305
3306	kv_indices are ragged global slot ids. With has_invalid they may hold
3307	-1, which the kernel masks; a caller whose stream never holds a
3308	negative slot passes False, which lets the fp8 kernel stage K straight
3309	into LDS. The cache format (bf16, per-tensor fp8, or DeepSeek V4's
3310	paged fp8_ds_mla) is inferred from kv_buffer and kv_scale. An fp8 q
3311	must come with its q_scale and runs both dots in fp8. The extra segment
3312	is DeepSeek V4's second cache: SWA window in kv_buffer, top-k
3313	compressed tokens here.
3314	"""
3315	from aiter.ops.triton.attention.sparse_mla import sparse_mla_fwd
3316
3317	fp8_q = q.dtype == FP8_DTYPE
3318	sparse_mla_fwd(
3319	q,
3320	kv_buffer,
3321	kv_indptr,
3322	kv_indices,
3323	sm_scale,
3324	kv_scale=kv_scale,
3325	kv_lora_rank=kv_lora_rank,
3326	qk_rope_head_dim=qk_rope_head_dim,
3327	has_invalid=has_invalid,
3328	dot_precision="fp8" if fp8_q else "bf16",
3329	q_scale=q_scale if fp8_q else None,
3330	out=o,
3331	attn_sink=attn_sink,
3332	extra_kv=extra_kv_buffer,
3333	extra_indptr=extra_kv_indptr,
3334	extra_indices=extra_kv_indices,
3335	)
3336
3337	@staticmethod
3338	def mla_decode_fwd(
3339	q: torch.Tensor,
3340	kv_buffer: torch.Tensor,
3341	o: torch.Tensor,
3342	sm_scale: float,
3343	qo_indptr: torch.Tensor,
3344	max_seqlen_qo: int,
3345	kv_indptr: torch.Tensor \| None = None,
3346	kv_indices: torch.Tensor \| None = None,
3347	kv_last_page_lens: torch.Tensor \| None = None,
3348	logit_cap: float = 0.0,
3349	q_scale: torch.Tensor \| None = None,
3350	kv_scale: torch.Tensor \| None = None,
3351	work_meta_data: torch.Tensor \| None = None,
3352	work_indptr: torch.Tensor \| None = None,
3353	work_info_set: torch.Tensor \| None = None,
3354	reduce_indptr: torch.Tensor \| None = None,
3355	reduce_final_map: torch.Tensor \| None = None,
3356	reduce_partial_map: torch.Tensor \| None = None,
3357	causal: bool = True,
3358	):
3359	torch.ops.vllm.rocm_aiter_mla_decode_fwd(
3360	q,
3361	kv_buffer.view(-1, 1, 1, q.shape[-1]),
3362	o,
3363	qo_indptr,
3364	max_seqlen_qo,
3365	kv_indptr,
3366	kv_indices,
3367	kv_last_page_lens,
3368	sm_scale=sm_scale,
3369	logit_cap=logit_cap,
3370	q_scale=q_scale,
3371	kv_scale=kv_scale,
3372	work_meta_data=work_meta_data,
3373	work_indptr=work_indptr,
3374	work_info_set=work_info_set,
3375	reduce_indptr=reduce_indptr,
3376	reduce_final_map=reduce_final_map,
3377	reduce_partial_map=reduce_partial_map,
3378	causal=causal,
3379	)
3380
3381	@staticmethod
3382	def mla_decode_fwd_lse(
3383	q: torch.Tensor,
3384	kv_buffer: torch.Tensor,
3385	o: torch.Tensor,
3386	sm_scale: float,
3387	qo_indptr: torch.Tensor,
3388	max_seqlen_qo: int,
3389	kv_indptr: torch.Tensor \| None = None,
3390	kv_indices: torch.Tensor \| None = None,
3391	kv_last_page_lens: torch.Tensor \| None = None,
3392	logit_cap: float = 0.0,
3393	q_scale: torch.Tensor \| None = None,
3394	kv_scale: torch.Tensor \| None = None,
3395	) -> torch.Tensor:
3396	"""Run MLA decode without persistent metadata and return its LSE."""
3397	return torch.ops.vllm.rocm_aiter_mla_decode_fwd_lse(
3398	q,
3399	kv_buffer.view(-1, 1, 1, q.shape[-1]),
3400	o,
3401	qo_indptr,
3402	max_seqlen_qo,
3403	kv_indptr,
3404	kv_indices,
3405	kv_last_page_lens,
3406	sm_scale=sm_scale,
3407	logit_cap=logit_cap,
3408	q_scale=q_scale,
3409	kv_scale=kv_scale,
3410	)
3411
3412	@staticmethod
3413	def per_tensor_quant(
3414	x: torch.Tensor,
3415	quant_dtype: torch.dtype,
3416	scale: torch.Tensor \| None = None,
3417	) -> tuple[torch.Tensor, torch.Tensor]:
3418	out = torch.empty_like(x, dtype=quant_dtype)
3419	is_dynamic = scale is None
3420	if is_dynamic:
3421	scale = torch.empty(1, dtype=torch.float32, device=x.device)
3422	torch.ops.vllm.rocm_aiter_per_tensor_quant(out, x, scale, is_dynamic)
3423	return out, scale
3424
3425	@staticmethod
3426	def per_token_quant(
3427	x: torch.Tensor,
3428	quant_dtype: torch.dtype,
3429	scale: torch.Tensor \| None = None,
3430	) -> tuple[torch.Tensor, torch.Tensor]:
3431	return torch.ops.vllm.rocm_aiter_per_token_quant(x, quant_dtype, scale)
3432
3433	@staticmethod
3434	def gemm_a8wfp4(
3435	x: torch.Tensor,
3436	w: torch.Tensor,
3437	x_scales: torch.Tensor,
3438	w_scales: torch.Tensor,
3439	out_dtype: torch.dtype,
3440	) -> torch.Tensor:
3441	return torch.ops.vllm.rocm_aiter_gemm_a8wfp4(
3442	x, w, x_scales, w_scales, out_dtype
3443	)
3444
3445	@staticmethod
3446	def triton_fp4_gemm_dynamic_quant(
3447	x: torch.Tensor,
3448	weight: torch.Tensor,
3449	weight_scale: torch.Tensor,
3450	out_dtype: torch.dtype \| None = torch.bfloat16,
3451	x_scales: torch.Tensor \| None = None,
3452	) -> torch.Tensor:
3453	from aiter.ops.triton.gemm_afp4wfp4 import gemm_afp4wfp4
3454	from aiter.ops.triton.quant import dynamic_mxfp4_quant
3455
3456	if x_scales is None:
3457	x_q, x_s = dynamic_mxfp4_quant(x)
3458	else:
3459	x_q = x
3460	x_s = x_scales
3461
3462	y = torch.empty(
3463	x_q.shape[0], weight.shape[0], device=x_q.device, dtype=out_dtype
3464	)
3465
3466	gemm_afp4wfp4(x_q, weight, x_s, weight_scale.T, out_dtype, y)
3467	return y
3468
3469	@staticmethod
3470	def fused_qk_norm_rope_and_cache(
3471	qkv: torch.Tensor,
3472	q_weight: torch.Tensor,
3473	k_weight: torch.Tensor,
3474	cos_sin_cache: torch.Tensor,
3475	positions: torch.Tensor,
3476	num_heads_q: int,
3477	num_heads_k: int,
3478	num_heads_v: int,
3479	head_dim: int,
3480	is_neox: bool,
3481	rms_norm_eps: float,
3482	q_out: torch.Tensor,
3483	k_cache: torch.Tensor,
3484	v_cache: torch.Tensor,
3485	slot_mapping: torch.Tensor,
3486	k_scale: torch.Tensor,
3487	v_scale: torch.Tensor,
3488	k_out: torch.Tensor \| None,
3489	v_out: torch.Tensor \| None,
3490	return_kv: bool,
3491	use_shuffle_layout: bool,
3492	block_size: int,
3493	x: int,
3494	rotary_dim: int = 0,
3495	):
3496	from aiter.ops.fused_qk_norm_rope_cache_quant import (
3497	fused_qk_norm_rope_cache_pts_quant_shuffle,
3498	)
3499
3500	fused_qk_norm_rope_cache_pts_quant_shuffle(
3501	qkv,
3502	q_weight,
3503	k_weight,
3504	cos_sin_cache,
3505	positions,
3506	qkv.size(0),
3507	num_heads_q,
3508	num_heads_k,
3509	num_heads_v,
3510	head_dim,
3511	is_neox,
3512	rms_norm_eps,
3513	q_out,
3514	k_cache,
3515	v_cache,
3516	slot_mapping,
3517	k_scale,
3518	v_scale,
3519	k_out,
3520	v_out,
3521	return_kv,
3522	use_shuffle_layout,
3523	block_size,
3524	x,
3525	rotary_dim,
3526	)
3527
3528	@staticmethod
3529	def do_qk_norm_rope_kvcache_update(
3530	qkv: torch.Tensor,
3531	q_weight: torch.Tensor,
3532	k_weight: torch.Tensor,
3533	cos_sin_cache: torch.Tensor,
3534	positions: torch.Tensor,
3535	num_heads_q: int,
3536	num_heads_k: int,
3537	head_dim: int,
3538	is_neox: bool,
3539	rms_norm_eps: float,
3540	q_out: torch.Tensor,
3541	k_out: torch.Tensor,
3542	key_cache: torch.Tensor,
3543	value_cache: torch.Tensor,
3544	slot_mapping: torch.Tensor,
3545	k_scale: torch.Tensor,
3546	v_scale: torch.Tensor,
3547	kv_cache_dtype: str,
3548	use_shuffle_layout: bool,
3549	) -> None:
3550	"""Run the fused QK-norm+RoPE+KV-cache op on already-split k/v caches.
3551
3552	Shared by the AITER FA and unified-attention impls. The caller splits
3553	kv_cache, since the unbind dim depends on the layout (e.g. the unified
3554	encoder-decoder path is K/V-first), and passes use_shuffle_layout
3555	(unified reads NHD and must pass False).
3556	"""
3557	if kv_cache_dtype.startswith("fp8"):
3558	key_cache = key_cache.view(current_platform.fp8_dtype())
3559	value_cache = value_cache.view(current_platform.fp8_dtype())
3560	# Partial-rotary support (e.g. GLM-4.7 applies rotary to only a prefix
3561	# of each head's channel dim).
3562	rotary_dim = cos_sin_cache.shape[-1]
3563	kernel_rotary_dim = 0 if rotary_dim == head_dim else rotary_dim
3564	rocm_aiter_ops.fused_qk_norm_rope_and_cache(
3565	qkv=qkv,
3566	q_weight=q_weight,
3567	k_weight=k_weight,
3568	cos_sin_cache=cos_sin_cache,
3569	positions=positions,
3570	num_heads_q=num_heads_q,
3571	num_heads_k=num_heads_k,
3572	num_heads_v=num_heads_k,
3573	head_dim=head_dim,
3574	is_neox=is_neox,
3575	rms_norm_eps=rms_norm_eps,
3576	q_out=q_out,
3577	k_cache=key_cache,
3578	v_cache=value_cache,
3579	slot_mapping=slot_mapping,
3580	k_scale=k_scale,
3581	v_scale=v_scale,
3582	k_out=k_out,
3583	v_out=None,
3584	return_kv=True,
3585	use_shuffle_layout=use_shuffle_layout,
3586	block_size=key_cache.shape[1],
3587	x=16 // key_cache.element_size(),
3588	rotary_dim=kernel_rotary_dim,
3589	)
3590
3591	@staticmethod
3592	def fused_qk_norm_mrope_and_cache(
3593	qkv: torch.Tensor,
3594	q_weight: torch.Tensor,
3595	k_weight: torch.Tensor,
3596	cos_sin_cache: torch.Tensor,
3597	positions: torch.Tensor,
3598	num_heads_q: int,
3599	num_heads_k: int,
3600	head_dim: int,
3601	is_neox: bool,
3602	mrope_section: list[int],
3603	is_interleaved: bool,
3604	rms_norm_eps: float,
3605	q_out: torch.Tensor,
3606	k_cache: torch.Tensor,
3607	v_cache: torch.Tensor,
3608	slot_mapping: torch.Tensor,
3609	k_scale: torch.Tensor,
3610	v_scale: torch.Tensor,
3611	block_size: int,
3612	x: int,
3613	rotary_dim: int = 0,
3614	) -> None:
3615	from aiter.ops.fused_qk_norm_mrope_cache_quant import (
3616	fused_qk_norm_mrope_3d_cache_pts_quant_shuffle,
3617	)
3618
3619	fused_qk_norm_mrope_3d_cache_pts_quant_shuffle(
3620	qkv,
3621	q_weight,
3622	k_weight,
3623	cos_sin_cache,
3624	positions,
3625	qkv.size(0),
3626	num_heads_q,
3627	num_heads_k,
3628	num_heads_k,
3629	head_dim,
3630	is_neox,
3631	mrope_section,
3632	is_interleaved,
3633	rms_norm_eps,
3634	q_out,
3635	k_cache,
3636	v_cache,
3637	slot_mapping,
3638	k_scale,
3639	v_scale,
3640	None,
3641	None,
3642	False,
3643	False,
3644	block_size,
3645	x,
3646	rotary_dim,
3647	)
3648
3649	@staticmethod
3650	def do_qk_norm_mrope_kvcache_update(
3651	qkv: torch.Tensor,
3652	q_weight: torch.Tensor,
3653	k_weight: torch.Tensor,
3654	cos_sin_cache: torch.Tensor,
3655	positions: torch.Tensor,
3656	num_heads_q: int,
3657	num_heads_k: int,
3658	head_dim: int,
3659	is_neox: bool,
3660	mrope_section: list[int],
3661	is_interleaved: bool,
3662	rms_norm_eps: float,
3663	q_out: torch.Tensor,
3664	key_cache: torch.Tensor,
3665	value_cache: torch.Tensor,
3666	slot_mapping: torch.Tensor,
3667	k_scale: torch.Tensor,
3668	v_scale: torch.Tensor,
3669	kv_cache_dtype: str,
3670	rotary_dim: int = 0,
3671	) -> None:
3672	if kv_cache_dtype.startswith("fp8"):
3673	key_cache = key_cache.view(current_platform.fp8_dtype())
3674	value_cache = value_cache.view(current_platform.fp8_dtype())
3675
3676	rocm_aiter_ops.fused_qk_norm_mrope_and_cache(
3677	qkv=qkv,
3678	q_weight=q_weight,
3679	k_weight=k_weight,
3680	cos_sin_cache=cos_sin_cache,
3681	positions=positions,
3682	num_heads_q=num_heads_q,
3683	num_heads_k=num_heads_k,
3684	head_dim=head_dim,
3685	is_neox=is_neox,
3686	mrope_section=mrope_section,
3687	is_interleaved=is_interleaved,
3688	rms_norm_eps=rms_norm_eps,
3689	q_out=q_out,
3690	k_cache=key_cache,
3691	v_cache=value_cache,
3692	slot_mapping=slot_mapping,
3693	k_scale=k_scale,
3694	v_scale=v_scale,
3695	block_size=key_cache.shape[1],
3696	x=16 // key_cache.element_size(),
3697	rotary_dim=rotary_dim,
3698	)
3699
3700	@staticmethod
3701	def triton_rope_and_cache(
3702	query: torch.Tensor,
3703	key: torch.Tensor,
3704	value: torch.Tensor,
3705	positions: torch.Tensor,
3706	cos_sin_cache: torch.Tensor,
3707	is_neox: bool,
3708	key_cache: torch.Tensor,
3709	value_cache: torch.Tensor,
3710	layer_slot_mapping: torch.Tensor,
3711	k_scale: torch.Tensor,
3712	v_scale: torch.Tensor,
3713	flash_layout: bool,
3714	apply_scale: bool,
3715	):
3716	from aiter.ops.triton.fused_kv_cache import fused_qk_rope_reshape_and_cache
3717
3718	cos, sin = cos_sin_cache.chunk(2, dim=-1)
3719	fused_qk_rope_reshape_and_cache(
3720	query,
3721	key,
3722	value,
3723	key_cache,
3724	value_cache,
3725	layer_slot_mapping,
3726	positions,
3727	cos,
3728	sin,
3729	k_scale,
3730	v_scale,
3731	is_neox,
3732	flash_layout=flash_layout,
3733	apply_scale=apply_scale,
3734	q_out=query,
3735	k_out=key,
3736	output_zeros=False,
3737	)
3738
3739	@staticmethod
3740	def batched_gemm_a16wfp4(
3741	X: torch.Tensor,
3742	W: torch.Tensor,
3743	w_scale: torch.Tensor,
3744	Y: torch.Tensor \| None = None,
3745	transpose_bm: bool \| None = False,
3746	prequant: bool \| None = False,
3747	y_scale: torch.Tensor \| None = None,
3748	) -> torch.Tensor:
3749	# ruff: noqa: E501 # isort: skip
3750	from aiter.ops.triton.batched_gemm_a16wfp4 import batched_gemm_a16wfp4
3751
3752	return batched_gemm_a16wfp4(
3753	X,
3754	W,
3755	w_scale,
3756	y=Y,
3757	transpose_bm=transpose_bm,
3758	prequant=prequant,
3759	y_scale=y_scale,
3760	)
3761
3762	@staticmethod
3763	def triton_fp8_bmm(
3764	X: torch.Tensor,
3765	WQ: torch.Tensor,
3766	w_scale: torch.Tensor,
3767	group_size: int = 128,
3768	bias: torch.Tensor \| None = None,
3769	dtype: torch.dtype \| None = torch.bfloat16,
3770	splitK: int \| None = None,
3771	YQ: torch.Tensor \| None = None,
3772	transpose_bm: bool \| None = False,
3773	config: dict \| None = None,
3774	) -> torch.Tensor:
3775	# ruff: noqa: E501 # isort: skip
3776	from aiter.ops.triton.batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant import (
3777	batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant as aiter_triton_fp8_bmm,
3778	)
3779
3780	return aiter_triton_fp8_bmm(
3781	X,
3782	WQ,
3783	w_scale,
3784	group_size=group_size,
3785	bias=bias,
3786	dtype=dtype,
3787	splitK=splitK,
3788	YQ=YQ,
3789	transpose_bm=transpose_bm,
3790	config=config,
3791	)
3792
3793	@staticmethod
3794	def group_fp8_quant(
3795	input_2d: torch.Tensor,
3796	group_size: int = 128,
3797	transpose_scale: bool = False,
3798	) -> tuple[torch.Tensor, torch.Tensor]:
3799	assert group_size == 128, "Group size must be 128"
3800	return torch.ops.vllm.rocm_aiter_group_fp8_quant(
3801	input_2d, group_size, transpose_scale
3802	)
3803
3804	@staticmethod
3805	@functools.cache
3806	def is_triton_gemm_w8a8_tuned(n: int, k: int) -> bool:
3807	if not current_platform.is_rocm():
3808	return False
3809	try:
3810	return _triton_gemm_config_is_tuned("GEMM-A8W8_BLOCKSCALE", n, k)
3811	except (AssertionError, ImportError):
3812	return False
3813
3814	@staticmethod
3815	@functools.cache
3816	def is_triton_gemm_afp4wfp4_presh_ws_tuned(n: int, k_bytes: int) -> bool:
3817	if not current_platform.is_rocm():
3818	return False
3819	# Input = weight.shape[1] (bytes); 4 for fp4 and aiter config 2k
3820	try:
3821	return _triton_gemm_config_is_tuned(
3822	"GEMM-AFP4WFP4_PRESHUFFLED", n, 4 * k_bytes
3823	)
3824	except (AssertionError, ImportError):
3825	return False
3826
3827	@staticmethod
3828	@functools.cache
3829	def is_shuffled_per_token_w8a8_gemm_tuned(
3830	N: int, K: int, q_dtype_w: torch.dtype
3831	) -> bool:
3832	return _ck_gemm_shape_is_tuned(
3833	N, K, q_dtype_w, "AITER_CONFIG_GEMM_A8W8_BPRESHUFFLE_FILE"
3834	)
3835
3836	@staticmethod
3837	@functools.cache
3838	def is_per_token_w8a8_gemm_tuned(N: int, K: int, q_dtype_w: torch.dtype) -> bool:
3839	return _ck_gemm_shape_is_tuned(N, K, q_dtype_w, "AITER_CONFIG_GEMM_A8W8_FILE")
3840
3841	@staticmethod
3842	def shuffle_weight(
3843	tensor: torch.Tensor, layout: tuple[int, int] = (16, 16)
3844	) -> torch.Tensor:
3845	from aiter.ops.shuffle import shuffle_weight
3846
3847	return shuffle_weight(tensor, layout=layout)
3848
3849	@staticmethod
3850	def shuffle_weight_a16w4(
3851	tensor: "torch.Tensor",
3852	nLane: int,
3853	gate_up: bool,
3854	) -> "torch.Tensor":
3855	"""Shuffles the weight tensor into (A16W4) layout for AITER kernels.
3856
3857	Args:
3858	tensor: The input weight tensor to be shuffled.
3859	nLane: Number of lanes in the shuffle layout.
3860	gate_up: Whether the weight is for w13 (True) or w2 (False).
3861
3862	Returns:
3863	torch.Tensor: The shuffled tensor.
3864
3865	"""
3866	from aiter.ops.shuffle import shuffle_weight_a16w4
3867
3868	return shuffle_weight_a16w4(tensor, nLane, gate_up)
3869
3870	@staticmethod
3871	def shuffle_scale_a16w4(
3872	tensor: "torch.Tensor",
3873	num_experts: int,
3874	gate_up: bool,
3875	) -> "torch.Tensor":
3876	"""Shuffles the scale tensor into (A16W4) layout for AITER kernels.
3877
3878	Args:
3879	tensor: The input scale tensor to be shuffled.
3880	num_experts: Number of experts, needed for reshaping logic.
3881	gate_up: Whether the scale is for w13 (True) or w2 (False).
3882
3883	Returns:
3884	torch.Tensor: The shuffled scale tensor.
3885
3886	"""
3887	from aiter.ops.shuffle import shuffle_scale_a16w4
3888
3889	return shuffle_scale_a16w4(tensor, num_experts, gate_up)
3890
3891	@staticmethod
3892	def shuffle_scale(tensor: torch.Tensor) -> torch.Tensor:
3893	"""Shuffle a non-gate/up E8M0 scale tensor for AITER."""
3894	from aiter.ops.shuffle import shuffle_scale
3895
3896	return shuffle_scale(tensor)
3897
3898	@staticmethod
3899	def shuffle_weights(
3900	*tensors: torch.Tensor, layout: tuple[int, int] = (16, 16)
3901	) -> tuple[torch.Tensor, ...]:
3902	"""Applies shuffle_weight function from AITER to each
3903	input tensor and returns them.
3904
3905	Rearranges (shuffles) the input tensor/s
3906	into a specified block layout for optimized computation.
3907
3908	Args:
3909	*tensors: Variable number of torch.Tensor objects.
3910	layout: A pair of integers specifying the block sizes used to divide
3911	the tensors during shuffling. Default is (16, 16).
3912
3913	Returns:
3914	A Tuple of shuffled tensors.
3915
3916	"""
3917	from aiter.ops.shuffle import shuffle_weight
3918
3919	return tuple(shuffle_weight(tensor, layout=layout) for tensor in tensors)
3920
3921	@staticmethod
3922	def shuffle_mxfp8_moe_weights(
3923	w13: torch.Tensor,
3924	w2: torch.Tensor,
3925	w13_scale: torch.Tensor,
3926	w2_scale: torch.Tensor,
3927	) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
3928	"""Preshuffle MXFP8 MoE weights + E8M0 scales into AITER's FlyDSL layout:
3929	gate/up-interleaved weights, interleaved scale for w13 (gate/up), plain
3930	scale for w2 (the interleaved variant is gate/up-only and misaligns w2).
3931	"""
3932	from aiter.ops.shuffle import shuffle_scale, shuffle_weight
3933
3934	num_experts = w13.shape[0]
3935	w13 = shuffle_weight(w13, is_guinterleave=True, gate_up=True)
3936	w2 = shuffle_weight(w2, is_guinterleave=True, gate_up=False)
3937	w13_scale = shuffle_scale(
3938	w13_scale.reshape(-1, w13_scale.shape[-1]),
3939	num_experts,
3940	is_guinterleave=True,
3941	gate_up=True,
3942	)
3943	w2_scale = shuffle_scale(w2_scale.reshape(-1, w2_scale.shape[-1]))
3944	return w13, w2, w13_scale, w2_scale
3945
3946	@staticmethod
3947	def flash_attn_varlen_func(
3948	q: torch.Tensor,
3949	k: torch.Tensor,
3950	v: torch.Tensor,
3951	cu_seqlens_q: torch.Tensor,
3952	cu_seqlens_k: torch.Tensor,
3953	max_seqlen_q: int,
3954	max_seqlen_k: int,
3955	min_seqlen_q: int \| None = None,
3956	dropout_p: float = 0.0,
3957	softmax_scale: float \| None = None,
3958	causal: bool = False,
3959	window_size: tuple[int, int] \| None = None,
3960	alibi_slopes: torch.Tensor \| None = None,
3961	return_lse: bool = False,
3962	out: torch.Tensor \| None = None,
3963	sink_ptr: torch.Tensor \| None = None,
3964	):
3965	"""Flash attention with variable length sequences.
3966
3967	This function is NOT wrapped with @is_aiter_supported decorator
3968	to allow explicit backend selection via attention_config to work
3969	even when VLLM_ROCM_USE_AITER=0.
3970
3971	Note: This performs lazy import of aiter.flash_attn_varlen_func
3972	"""
3973	from aiter import flash_attn_varlen_func
3974
3975	return flash_attn_varlen_func(
3976	q=q,
3977	k=k,
3978	v=v,
3979	cu_seqlens_q=cu_seqlens_q,
3980	cu_seqlens_k=cu_seqlens_k,
3981	max_seqlen_q=max_seqlen_q,
3982	max_seqlen_k=max_seqlen_k,
3983	min_seqlen_q=min_seqlen_q,
3984	dropout_p=dropout_p,
3985	softmax_scale=softmax_scale,
3986	causal=causal,
3987	window_size=window_size,
3988	alibi_slopes=alibi_slopes,
3989	return_lse=return_lse,
3990	out=out,
3991	sink_ptr=sink_ptr,
3992	)
3993
3994	@staticmethod
3995	def fp8_attn_wrapper(
3996	q: torch.Tensor,
3997	k: torch.Tensor,
3998	v: torch.Tensor,
3999	q_descale: torch.Tensor,
4000	k_descale: torch.Tensor,
4001	v_descale: torch.Tensor,
4002	batch_size: int,
4003	output_dtype: torch.dtype,
4004	scale: float \| None = None,
4005	cu_seqlens: torch.Tensor \| None = None,
4006	max_seqlen: torch.Tensor \| None = None,
4007	) -> torch.Tensor:
4008	return torch.ops.vllm.aiter_fp8_attn_wrapper(
4009	q,
4010	k,
4011	v,
4012	q_descale,
4013	k_descale,
4014	v_descale,
4015	batch_size,
4016	output_dtype,
4017	scale,
4018	cu_seqlens,
4019	max_seqlen,
4020	)
4021
4022	@staticmethod
4023	def pa_fwd_asm(
4024	Q: torch.Tensor,
4025	K: torch.Tensor,
4026	V: torch.Tensor,
4027	block_tables: torch.Tensor,
4028	context_lens: torch.Tensor,
4029	block_tables_stride0: int,
4030	K_QScale: torch.Tensor,
4031	V_QScale: torch.Tensor,
4032	out_: torch.Tensor,
4033	):
4034	"""Paged attention forward pass using assembly kernel.
4035
4036	This function is NOT wrapped with @is_aiter_supported decorator
4037	to allow explicit backend selection via attention_config to work
4038	even when VLLM_ROCM_USE_AITER=0.
4039
4040	Note: This performs lazy import of aiter.pa_fwd_asm
4041	"""
4042	from aiter import pa_fwd_asm
4043
4044	return pa_fwd_asm(
4045	Q=Q,
4046	K=K,
4047	V=V,
4048	block_tables=block_tables,
4049	context_lens=context_lens,
4050	block_tables_stride0=block_tables_stride0,
4051	K_QScale=K_QScale,
4052	V_QScale=V_QScale,
4053	out_=out_,
4054	)
4055
4056	@staticmethod
4057	def paged_attention_common(
4058	Q: torch.Tensor,
4059	K: torch.Tensor,
4060	V: torch.Tensor,
4061	tmp_out: torch.Tensor,
4062	max_logits: torch.Tensor,
4063	exp_sums: torch.Tensor,
4064	max_seq_len: int,
4065	block_tables: torch.Tensor,
4066	context_lens: torch.Tensor,
4067	block_tables_stride0: int,
4068	scale: float,
4069	K_QScale_hip: torch.Tensor,
4070	V_QScale_hip: torch.Tensor,
4071	K_QScale_asm: torch.Tensor,
4072	V_QScale_asm: torch.Tensor,
4073	out_: torch.Tensor,
4074	kv_cache_dtype: str,
4075	):
4076	"""Paged attention common function.
4077
4078	This function is NOT wrapped with @is_aiter_supported decorator
4079	to allow explicit backend selection via attention_config to work
4080	even when VLLM_ROCM_USE_AITER=0.
4081
4082	Note: This performs lazy import of aiter.paged_attention_common
4083	"""
4084	from aiter import paged_attention_common
4085
4086	return paged_attention_common(
4087	Q=Q,
4088	K=K,
4089	V=V,
4090	tmp_out=tmp_out,
4091	max_logits=max_logits,
4092	exp_sums=exp_sums,
4093	max_seq_len=max_seq_len,
4094	block_tables=block_tables,
4095	context_lens=context_lens,
4096	block_tables_stride0=block_tables_stride0,
4097	scale=scale,
4098	K_QScale_hip=K_QScale_hip,
4099	V_QScale_hip=V_QScale_hip,
4100	K_QScale_asm=K_QScale_asm,
4101	V_QScale_asm=V_QScale_asm,
4102	out_=out_,
4103	kv_cache_dtype=kv_cache_dtype,
4104	)
4105
4106	@staticmethod
4107	def mhc_pre(
4108	residual: torch.Tensor,
4109	fn: torch.Tensor,
4110	hc_scale: torch.Tensor,
4111	hc_base: torch.Tensor,
4112	rms_eps: float,
4113	hc_pre_eps: float,
4114	hc_sinkhorn_eps: float,
4115	hc_post_mult_value: float,
4116	sinkhorn_repeat: int,
4117	norm_weight: torch.Tensor \| None = None,
4118	norm_eps: float = 0.0,
4119	) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
4120	"""Forward pass for mHC pre block.
4121
4122	Args:
4123	residual: shape (..., hc_mult, hidden_size), dtype torch.bfloat16
4124	fn: shape (hc_mult3, hc_mult * hidden_size), dtype torch.float32
4125	hc_scale: shape (3,), dtype torch.float32
4126	hc_base: shape (hc_mult3,), dtype torch.float32
4127	rms_eps: RMS normalization epsilon
4128	hc_pre_eps: pre-mix epsilon
4129	hc_sinkhorn_eps: sinkhorn epsilon
4130	hc_post_mult_value: post-mix multiplier value
4131	sinkhorn_repeat: number of sinkhorn iterations
4132	norm_weight: optional RMSNorm weight fused into the pre kernel
4133	norm_eps: epsilon for the fused RMSNorm when norm_weight is set
4134
4135	Returns:
4136	post_mix: shape (..., hc_mult), dtype torch.float32
4137	comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32
4138	layer_input: shape (..., hidden_size), dtype torch.bfloat16
4139
4140	"""
4141	from aiter.ops.mhc import mhc_pre
4142
4143	# Validate shapes
4144	assert residual.dtype == torch.bfloat16
4145	assert fn.dtype == torch.float32
4146	assert hc_scale.dtype == torch.float32
4147	assert hc_base.dtype == torch.float32
4148
4149	hc_mult = residual.shape[-2]
4150	hidden_size = residual.shape[-1]
4151	hc_mult2 = hc_mult * hc_mult
4152	hc_mult3 = hc_mult * 2 + hc_mult2
4153
4154	hc_hidden_size = hc_mult * hidden_size
4155	assert fn.shape[0] == hc_mult3
4156	assert fn.shape[1] == hc_hidden_size
4157	assert hc_scale.shape == (3,)
4158	assert hc_base.shape == (hc_mult3,)
4159
4160	outer_shape = residual.shape[:-2]
4161
4162	residual_flat = residual.view(-1, hc_mult, hidden_size)
4163
4164	num_tokens = residual_flat.shape[0]
4165	if num_tokens == 0:
4166	return (
4167	torch.empty(
4168	num_tokens,
4169	hc_mult,
4170	1,
4171	dtype=torch.float32,
4172	device=residual_flat.device,
4173	),
4174	torch.empty(
4175	num_tokens,
4176	hc_mult,
4177	hc_mult,
4178	dtype=torch.float32,
4179	device=residual_flat.device,
4180	),
4181	torch.empty(
4182	num_tokens,
4183	hidden_size,
4184	dtype=torch.bfloat16,
4185	device=residual_flat.device,
4186	),
4187	)
4188
4189	# AITER's Python wrapper allocates intermediate/output tensors without
4190	# explicit device arguments, so run it under the residual tensor's device.
4191	with torch.device(residual_flat.device):
4192	args = (
4193	residual_flat,
4194	fn,
4195	hc_scale,
4196	hc_base,
4197	rms_eps,
4198	hc_pre_eps,
4199	hc_sinkhorn_eps,
4200	hc_post_mult_value,
4201	sinkhorn_repeat,
4202	)
4203	if norm_weight is None:
4204	post_mix, comb_mix, layer_input = mhc_pre(*args)
4205	else:
4206	post_mix, comb_mix, layer_input = mhc_pre(*args, norm_weight, norm_eps)
4207	return (
4208	post_mix.view(*outer_shape, hc_mult, 1),
4209	comb_mix.view(*outer_shape, hc_mult, hc_mult),
4210	layer_input.view(*outer_shape, hidden_size),
4211	)
4212
4213	@staticmethod
4214	def mhc_pre_delayed(
4215	residual: torch.Tensor,
4216	fn: torch.Tensor,
4217	hc_scale: torch.Tensor,
4218	hc_base: torch.Tensor,
4219	rms_eps: float,
4220	hc_pre_eps: float,
4221	hc_sinkhorn_eps: float,
4222	hc_post_mult_value: float,
4223	sinkhorn_repeat: int,
4224	pre_mix: torch.Tensor \| None = None,
4225	sublayer_out: torch.Tensor \| None = None,
4226	post_layer_mix: torch.Tensor \| None = None,
4227	comb_res_mix: torch.Tensor \| None = None,
4228	residual_out: torch.Tensor \| None = None,
4229	) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
4230	"""MHC pre using the pre-mix carried from the previous sublayer.
4231
4232	Same gates as :meth:`mhc_pre`, but the stream collapse applies the
4233	caller's ``pre_mix`` instead of the one computed here, and the one
4234	computed here is returned for the next sublayer seam. AITER has no
4235	single kernel for that shape, so drive its two stages directly:
4236	``mhc_pre_big_fuse`` still produces the post and comb gates (including
4237	every sinkhorn iteration), while the pre gate is recovered from the
4238	same split-k GEMM output and the collapse is done by the Triton
4239	kernel. ``mhc_pre_big_fuse`` also writes a collapse we do not use;
4240	that redundant store is the price of not having a native delayed
4241	kernel, and is small next to the ~140 launches it replaces.
4242
4243	Passing ``sublayer_out`` / ``post_layer_mix`` / ``comb_res_mix``
4244	also applies the preceding post block. AITER then projects the new
4245	residual in the same kernel that computes it
4246	(``mhc_fused_post_pre_gemm_sqrsum``), which saves one launch; every
4247	later stage is identical either way. That new residual is written to
4248	``residual_out``, which the caller must supply in that case; it is an
4249	output buffer rather than a return value so the op never has to hand
4250	back one of its own inputs on the unfused path.
4251
4252	Returns:
4253	post_mix: shape (..., hc_mult, 1), dtype torch.float32
4254	comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32
4255	layer_input: shape (..., hidden_size), dtype torch.bfloat16
4256	next_pre_mix: shape (..., hc_mult), dtype torch.float32
4257
4258	"""
4259	from aiter.ops.mhc import (
4260	get_mhc_fused_post_pre_config,
4261	get_mhc_pre_splitk,
4262	mhc_fused_post_pre_gemm_sqrsum,
4263	mhc_pre_gemm_sqrsum,
4264	)
4265
4266	assert residual.dtype == torch.bfloat16
4267	assert fn.dtype == torch.float32
4268	assert hc_scale.dtype == torch.float32
4269	assert hc_base.dtype == torch.float32
4270
4271	hc_mult = residual.shape[-2]
4272	hidden_size = residual.shape[-1]
4273	hc_mult3 = hc_mult * 2 + hc_mult * hc_mult
4274	hc_hidden_size = hc_mult * hidden_size
4275
4276	assert fn.shape == (hc_mult3, hc_hidden_size)
4277	assert hc_scale.shape == (3,)
4278	assert hc_base.shape == (hc_mult3,)
4279
4280	outer_shape = residual.shape[:-2]
4281	residual_flat = residual.view(-1, hc_mult, hidden_size)
4282	num_tokens = residual_flat.shape[0]
4283	device = residual_flat.device
4284
4285	if num_tokens == 0:
4286	return (
4287	torch.empty(0, hc_mult, 1, dtype=torch.float32, device=device),
4288	torch.empty(0, hc_mult, hc_mult, dtype=torch.float32, device=device),
4289	torch.empty(0, hidden_size, dtype=torch.bfloat16, device=device),
4290	torch.empty(0, hc_mult, dtype=torch.float32, device=device),
4291	)
4292
4293	pre_mix_flat = None if pre_mix is None else pre_mix.view(-1, hc_mult)
4294
4295	fuse_post = sublayer_out is not None
4296	# AITER's Python wrappers allocate without explicit device arguments.
4297	with torch.device(device):
4298	if fuse_post:
4299	# Keyed on the per-stream K, i.e. hidden_size, not hc_mult * it.
4300	splitk, tile_m, tile_n, tile_k = get_mhc_fused_post_pre_config(
4301	num_tokens, hidden_size
4302	)
4303	else:
4304	splitk, tile_k = get_mhc_pre_splitk(num_tokens, hc_hidden_size)
4305	# AITER pads the GEMM output to a multiple of 32 columns.
4306	gemm_pad = torch.empty(
4307	splitk,
4308	num_tokens,
4309	(hc_mult3 + 31) // 32 * 32,
4310	dtype=torch.float32,
4311	device=device,
4312	)
4313	gemm_out = gemm_pad[:, :, :hc_mult3]
4314	sqrsum = torch.empty(splitk, num_tokens, dtype=torch.float32, device=device)
4315	if fuse_post:
4316	# Restated for the type checker, which cannot narrow the
4317	# optional arguments through the fuse_post flag.
4318	assert sublayer_out is not None
4319	assert post_layer_mix is not None and comb_res_mix is not None
4320	assert residual_out is not None
4321	projected = residual_out.view(-1, hc_mult, hidden_size)
4322	mhc_fused_post_pre_gemm_sqrsum(
4323	gemm_out,
4324	sqrsum,
4325	projected,
4326	sublayer_out.view(-1, hidden_size),
4327	residual_flat,
4328	# The kernel wants the post mix squeezed to (tokens, hc_mult).
4329	post_layer_mix.view(-1, hc_mult, 1).squeeze(-1),
4330	comb_res_mix.view(-1, hc_mult, hc_mult),
4331	fn,
4332	tile_m,
4333	tile_n,
4334	tile_k,
4335	0,
4336	)
4337	else:
4338	projected = residual_flat
4339	mhc_pre_gemm_sqrsum(gemm_out, sqrsum, projected, fn, tile_k, 0)
4340
4341	post_mix, comb_mix, layer_input, next_pre_mix = _mhc_delayed_pre_tail(
4342	gemm_out,
4343	sqrsum,
4344	projected,
4345	hc_scale,
4346	hc_base,
4347	pre_mix_flat,
4348	rms_eps,
4349	hc_pre_eps,
4350	hc_sinkhorn_eps,
4351	hc_post_mult_value,
4352	sinkhorn_repeat,
4353	)
4354
4355	return (
4356	post_mix.view(*outer_shape, hc_mult, 1),
4357	comb_mix.view(*outer_shape, hc_mult, hc_mult),
4358	layer_input.view(*outer_shape, hidden_size),
4359	next_pre_mix.view(*outer_shape, hc_mult),
4360	)
4361
4362	@staticmethod
4363	def mhc_fused_post_pre_delayed_rms_norm(
4364	residual: torch.Tensor,
4365	fn: torch.Tensor,
4366	hc_scale: torch.Tensor,
4367	hc_base: torch.Tensor,
4368	rms_eps: float,
4369	hc_pre_eps: float,
4370	hc_sinkhorn_eps: float,
4371	hc_post_mult_value: float,
4372	sinkhorn_repeat: int,
4373	pre_mix: torch.Tensor \| None,
4374	sublayer_out: torch.Tensor \| None,
4375	post_layer_mix: torch.Tensor \| None,
4376	comb_res_mix: torch.Tensor \| None,
4377	norm_weight: torch.Tensor,
4378	norm_eps: float,
4379	residual_out: torch.Tensor \| None,
4380	) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
4381	""":meth:`mhc_pre_delayed` with the RMSNorm of the collapse folded in.
4382
4383	Runs aiter's fused Triton kernel ``mhc_fused_post_pre_delayed_rmsnorm``;
4384	the returned ``layer_input`` is already normalised with ``norm_weight`` /
4385	``norm_eps``.
4386	"""
4387	from aiter.ops.triton.fusions.mhc_fused_post_pre_delayed_rmsnorm import (
4388	mhc_fused_post_pre_delayed_rmsnorm,
4389	)
4390
4391	_, post_mix, comb_mix, layer_input, next_pre_mix = (
4392	mhc_fused_post_pre_delayed_rmsnorm(
4393	residual,
4394	fn,
4395	hc_scale,
4396	hc_base,
4397	rms_eps,
4398	hc_pre_eps,
4399	hc_sinkhorn_eps,
4400	hc_post_mult_value,
4401	sinkhorn_repeat,
4402	pre_mix,
4403	sublayer_out,
4404	post_layer_mix,
4405	comb_res_mix,
4406	norm_weight,
4407	norm_eps,
4408	residual_out=residual_out,
4409	)
4410	)
4411	return post_mix, comb_mix, layer_input, next_pre_mix
4412
4413	@staticmethod
4414	def mhc_fused_post_pre_delayed_prefers_unfused(num_tokens: int) -> bool:
4415	"""True when AITER's own heuristic favours separate post and pre.
4416
4417	Mirrors the ``fused_m_upper_bound`` table in ``aiter.ops.mhc``: the
4418	fused GEMM wins for decode-sized batches and loses for large ones.
4419	"""
4420	from aiter.jit.utils.chip_info import get_gfx_runtime
4421
4422	return num_tokens >= {
4423	"gfx950": 1024,
4424	"gfx942": 128,
4425	"gfx1250": 1024,
4426	}.get(get_gfx_runtime(), 1024)
4427
4428	@staticmethod
4429	def hc_head(
4430	hs_flat: torch.Tensor,
4431	fn: torch.Tensor,
4432	hc_scale: torch.Tensor,
4433	hc_base: torch.Tensor,
4434	out: torch.Tensor,
4435	hidden_size: int,
4436	rms_eps: float,
4437	hc_eps: float,
4438	hc_mult: int,
4439	) -> None:
4440	"""Run hc_head through AITER mhc_pre and write the result to out."""
4441	assert hs_flat.dtype == torch.bfloat16
4442	assert fn.dtype == torch.float32
4443	assert hc_scale.dtype == torch.float32
4444	assert hc_base.dtype == torch.float32
4445	assert hs_flat.shape[-2:] == (hc_mult, hidden_size)
4446	assert fn.shape == (hc_mult, hc_mult * hidden_size)
4447	assert hc_scale.shape == (1,)
4448	assert hc_base.shape == (hc_mult,)
4449
4450	num_tokens = hs_flat.shape[0]
4451	if num_tokens == 0:
4452	return
4453
4454	hc_mult3 = hc_mult * 2 + hc_mult * hc_mult
4455
4456	full_fn = torch.zeros(
4457	hc_mult3,
4458	hc_mult * hidden_size,
4459	dtype=fn.dtype,
4460	device=fn.device,
4461	)
4462	full_fn[:hc_mult] = fn
4463
4464	full_base = torch.zeros(hc_mult3, dtype=hc_base.dtype, device=hc_base.device)
4465	full_base[:hc_mult] = hc_base
4466
4467	full_scale = torch.zeros(3, dtype=hc_scale.dtype, device=hc_scale.device)
4468	full_scale[0] = hc_scale[0]
4469
4470	_, _, layer_input = rocm_aiter_ops.mhc_pre(
4471	hs_flat,
4472	full_fn,
4473	full_scale,
4474	full_base,
4475	rms_eps,
4476	hc_eps,
4477	0.0,
4478	1.0,
4479	0,
4480	)
4481	out.copy_(layer_input)
4482
4483	@staticmethod
4484	def mhc_post(
4485	x: torch.Tensor,
4486	residual: torch.Tensor,
4487	post_layer_mix: torch.Tensor,
4488	comb_res_mix: torch.Tensor,
4489	) -> torch.Tensor:
4490	from aiter.ops.mhc import mhc_post
4491
4492	hc_mult = residual.shape[-2]
4493	hidden_size = residual.shape[-1]
4494	residual_flat = residual.view(-1, hc_mult, hidden_size)
4495	num_tokens = residual_flat.shape[0]
4496	out = torch.empty_like(residual_flat)
4497	mhc_post(
4498	out,
4499	x.view(num_tokens, hidden_size),
4500	residual_flat,
4501	post_layer_mix.view(num_tokens, hc_mult, 1),
4502	comb_res_mix.view(num_tokens, hc_mult, hc_mult),
4503	)
4504	return out.view_as(residual)
4505
4506	@staticmethod
4507	def mhc_fused_post_pre(
4508	x: torch.Tensor,
4509	residual: torch.Tensor,
4510	post_layer_mix: torch.Tensor,
4511	comb_res_mix: torch.Tensor,
4512	fn: torch.Tensor,
4513	hc_scale: torch.Tensor,
4514	hc_base: torch.Tensor,
4515	rms_eps: float,
4516	hc_pre_eps: float,
4517	hc_sinkhorn_eps: float,
4518	hc_post_mult_value: float,
4519	sinkhorn_repeat: int,
4520	norm_weight: torch.Tensor \| None = None,
4521	norm_eps: float = 0.0,
4522	) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
4523	"""Fused mHC post + next mHC pre via AITER.
4524
4525	Returns residual_cur, post_mix, comb_mix, layer_input in vLLM order
4526	(AITER returns post_mix, comb_mix, layer_input, next_residual).
4527	"""
4528	from aiter.ops.mhc import mhc_fused_post_pre
4529
4530	assert x.dtype == torch.bfloat16
4531	assert residual.dtype == torch.bfloat16
4532	assert fn.dtype == torch.float32
4533	assert hc_scale.dtype == torch.float32
4534	assert hc_base.dtype == torch.float32
4535
4536	hc_mult = residual.shape[-2]
4537	hidden_size = residual.shape[-1]
4538	outer_shape = residual.shape[:-2]
4539
4540	residual_flat = residual.view(-1, hc_mult, hidden_size)
4541	num_tokens = residual_flat.shape[0]
4542	x_flat = x.view(num_tokens, hidden_size)
4543	post_flat = post_layer_mix.view(num_tokens, hc_mult, 1)
4544	comb_flat = comb_res_mix.view(num_tokens, hc_mult, hc_mult)
4545
4546	if num_tokens == 0:
4547	return (
4548	torch.empty_like(residual_flat).view_as(residual),
4549	torch.empty(
4550	*outer_shape,
4551	hc_mult,
4552	1,
4553	dtype=torch.float32,
4554	device=residual.device,
4555	),
4556	torch.empty(
4557	*outer_shape,
4558	hc_mult,
4559	hc_mult,
4560	dtype=torch.float32,
4561	device=residual.device,
4562	),
4563	torch.empty(
4564	*outer_shape,
4565	hidden_size,
4566	dtype=torch.bfloat16,
4567	device=residual.device,
4568	),
4569	)
4570
4571	args = (
4572	x_flat,
4573	residual_flat,
4574	post_flat,
4575	comb_flat,
4576	fn,
4577	hc_scale,
4578	hc_base,
4579	rms_eps,
4580	hc_pre_eps,
4581	hc_sinkhorn_eps,
4582	hc_post_mult_value,
4583	sinkhorn_repeat,
4584	)
4585	with torch.device(residual_flat.device):
4586	if norm_weight is None:
4587	result = mhc_fused_post_pre(*args)
4588	else:
4589	result = mhc_fused_post_pre(*args, norm_weight, norm_eps)
4590	post_mix, comb_mix, layer_input, next_residual = result
4591
4592	return (
4593	next_residual.view_as(residual),
4594	post_mix.view(*outer_shape, hc_mult, 1),
4595	comb_mix.view(*outer_shape, hc_mult, hc_mult),
4596	layer_input.view(*outer_shape, hidden_size),
4597	)
4598
4599
4600	rocm_aiter_ops.register_ops_once()
4601	_sync_aiter_situv2_moe_env()