[WIP] SP+AsyncTP piecewise compilation fix + per-matmul heuristic gating

[tianrengao]

Summary

Fix SP+AsyncTP crash with piecewise compilation and add per-matmul heuristic gating.

Changes

1. Piecewise crash fix: _remove_mismatched_epilogue_copies() + _fix_cross_submod_outputs() — SP changes tensor shapes at submod boundaries, causing assert_size_stride failures
2. Per-matmul heuristic: Decision tree from AutoHeuristic gates fusion via extra_check, preventing 2-3x decode regression at small batch sizes
3. Config: async_tp_min_tokens threshold + enable_sp_and_async_tp() runtime check

Status

⚠️ Draft — SP is currently disabled in upstream (IS_DENSE=False, #25689). These changes are preserved for future work.

Benchmark results show no meaningful serving benefit on single-node NVLink:

• Llama-2-7b TP=2: 1.01-1.02x (noise)
• Meta-Llama-3-70B TP=8: 0.98x

See benchmarks/SP_ASYNCTP_BENCHMARK_REPORT.md for details.

[gemini-code-assist]

Code Review

This pull request introduces a fix for a crash when using Sequence Parallelism (SP) with Asynchronous Tensor Parallelism (AsyncTP) in piecewise compilation mode. It also adds a per-matmul heuristic to gate fusions, preventing performance regressions at small batch sizes. The changes include new logic to handle tensor shape modifications at submodule boundaries, a new configuration option async_tp_min_tokens for coarse-grained gating, and new benchmark scripts and tests. My review found one high-severity issue related to error handling in the new heuristic check. Otherwise, the changes appear well-structured and address the described problems effectively.

[gemini-code-assist]

The use of a broad except Exception: is risky as it can hide various errors during shape extraction for the heuristic check (e.g., KeyError, AttributeError, IndexError). If an unexpected error occurs, it will be silently caught, and the function will default to returning True, causing the fusion to be applied. This could lead to incorrect behavior or mask underlying bugs in the pattern matching. It would be safer to catch more specific exceptions and log a warning to facilitate debugging.

Suggested change

	except Exception:
	return True
	except (KeyError, IndexError, AttributeError) as e:
	logger.warning(
	"Failed to extract matmul shapes for async TP heuristic check: %s. "
	"Defaulting to allowing fusion.",
	e,
	exc_info=True,
	)
	return True

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @tianrengao.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @tianrengao.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[WoosukKwon]

Closing as superseded by #60517, which removed the torch.compile-based sequence parallelism pass and AsyncTP (along with enable_sp,
fuse_gemm_comms, sp_min_token_num, and their tests). The code this PR modifies no longer exists. Thanks for the contribution!