[ROCm] AITER fused allreduce+rmsnorm: merge PRs #37646 + #81 with reviewer fixes

[rbrugaro-amd]

…iewer fixes

Merge #37646 (AITER fused allreduce + RMSNorm pass) and
EmbeddedLLM#81 (hidden_dim=7168 support + threshold tuning), then
address maintainer feedback:

• Add AITER compile range splitting in _set_compile_ranges() so the backend generates separate compiled graphs for fusion-eligible vs non-eligible token counts, matching the existing flashinfer pattern.
• Convert RocmAiterAllReduceFusionPass to VllmFusionPatternMatcherPass and rewrite AITER patterns as VllmPatternReplacement subclasses. Add extra_check support to VllmPatternReplacement/register().
• Extract _compute_allreduce_max_token_num() shared helper used by both AllReduceFusionPass (flashinfer) and RocmAiterAllReduceFusionPass.
• Simplify parallel_state.py aiter capture context per review suggestion.

[gemini-code-assist]

Code Review

This pull request implements fused All-Reduce and RMSNorm operations for ROCm platforms using the AITER library, including a new compilation pass (RocmAiterAllReduceFusionPass) and integration with the distributed graph capture. Feedback identifies a bug where the residual buffer in the fused pattern must be initialized to zero to avoid incorrect results from garbage data. Additionally, the lifecycle management of the AITER communicator needs correction, as destroying it in the pass's destructor causes runtime crashes during execution, and its initialization should be made idempotent.

[gemini-code-assist]

The AiterAllreduceFusedRMSNormPattern replacement uses torch.empty_like(input) for the residual buffer. Since the underlying AITER kernel performs an additive operation (allreduce(input) + residual), using an uninitialized buffer will lead to incorrect results containing garbage data. This buffer should be initialized to zeros for the non-additive case to ensure correctness.

Suggested change

	residual = torch.empty_like(input)
	allreduce = op(
	input_=input,
	residual=residual,
	weight=weight,
	epsilon=epsilon,
	)
	return allreduce[0], allreduce[1]
	residual = torch.zeros_like(input)
	allreduce = op(
	input_=input,
	residual=residual,
	weight=weight,
	epsilon=epsilon,
	)
	return allreduce[0], allreduce[1]

[gemini-code-assist]

The __del__ method destroys the global AITER allreduce communicator. This is highly problematic because the pass instance is typically garbage collected after the compilation phase, but the global communicator must remain alive for the execution of the compiled graph. Destroying it here will cause runtime crashes when the model is executed. Furthermore, in environments with multiple engine instances or models, deleting one pass will break the communicator for all others. The lifecycle of the communicator should be managed by a persistent entity like parallel_state or GroupCoordinator, not by a fusion pass instance.

[gemini-code-assist]

The initialize_aiter_allreduce method should check if the communicator is already initialized before creating a new one. This prevents redundant resource allocation and potential race conditions when multiple fusion passes are instantiated (e.g., during multiple compilation cycles or for different models in the same process).

Suggested change

	try:
	from aiter.dist.device_communicators.custom_all_reduce import (
	CustomAllreduce as AiterCustomAllreduce,
	)
	cls._CUSTOM_ALL_REDUCE = AiterCustomAllreduce(group, device)
	except Exception:
	cls._CUSTOM_ALL_REDUCE = None
	if cls._CUSTOM_ALL_REDUCE is not None:
	return
	try:
	from aiter.dist.device_communicators.custom_all_reduce import (
	CustomAllreduce as AiterCustomAllreduce,
	)
	cls._CUSTOM_ALL_REDUCE = AiterCustomAllreduce(group, device)
	except Exception:
	cls._CUSTOM_ALL_REDUCE = None