CP + SDPA issue w/ SAC when using TorchFunctionMode

[XilunWu]

Bug description

Context Parallel has 2 ways to dispatch SDPA to the corresponding KV all-gather variant: 1. through monkey-patch; 2. using TorchFunctionMode. The first approach works well with SAC but the second is failing with error "RuntimeError: view size is not compatible with input tensor's size and stride (at least one dimension spans across two contiguous subspaces). Use .reshape(...) instead."

See reproduce PR: #1434

Versions

source build: 32aade9d8d3