Support context-parallel local_map in AutoParallel

[AlbedoWang]

Adds Context Parallel attention support through local_map and teaches AutoParallel's local_map HOP path to honor in_grad_placements.

Context Parallel needs asymmetric placements for attention inputs and their gradients. Q remains sequence-sharded on the CP axis in both forward and backward, while K/V are replicated for the forward attention kernel and produce partial gradients in backward:

• Q forward/grad: Shard(sequence) on CP
• K/V forward: Replicate() on CP
• K/V grad: Partial() on CP

This matches Torchtitan full-DTensor CP semantics. The gap was specific to AutoParallel's tracing path: AutoParallel enables PyTorch's local_map_hop wrapping so local_map regions remain explicit placement-aware graph boundaries. That HOP autograd path currently rejects in_grad_placements, and its generated backward metadata defaults backward output placements to the forward input placements. That is correct for existing EP/MoE local_map regions with in_grad_placements=None, but it is wrong for CP K/V gradients.

This PR adds a scoped AutoParallel shim around local_map HOP autograd:

• bypass the upstream in_grad_placements reject while AutoParallel is tracing or applying placements
• keep PyTorch's existing HOP forward/backward graph construction
• rewrite backward local_map out_placements to the forward in_grad_placements when provided
• restore the original PyTorch HOP implementation when leaving the context

The placement optimizer no longer rejects in_grad_placements; it continues to consume local_map in_placements and out_placements. For backward HOPs, the shim has already made out_placements reflect the gradient placements, so CP logic stays out of the generic placement-option path.

Also adds autoparallel.context_parallel helpers for CP attention local_map wrapping, including an SDPA convenience wrapper. The helpers support named 2D, 3D, and 4D meshes such as:

• ("dp_shard", "cp")
• ("dp_shard", "tp")
• ("dp_shard", "cp", "tp")
• ("dp_replicate", "dp_shard", "cp", "tp")

Docs included:

• docs/context_parallel.md: user-facing usage and placement semantics
• docs/context_parallel_design.md: implementation rationale and unsupported path analysis

Review order:

1. autoparallel/tracing.py — local_map HOP autograd shim and scoped enablement
2. autoparallel/shardings/placement_options.py — local_map placement option support for in-grad metadata
3. autoparallel/context_parallel.py — CP placement helper and SDPA wrapper
4. tests/test_context_parallel.py and tests/test_correctness.py — end-to-end coverage and shared correctness helper extension
5. docs/context_parallel.md and docs/context_parallel_design.md

Test plan:

• pre-commit run --files autoparallel/api.py autoparallel/shardings/placement_options.py autoparallel/tracing.py autoparallel/context_parallel.py tests/test_correctness.py tests/test_context_parallel.py docs/README.md docs/context_parallel.md docs/context_parallel_design.md
• git diff --check
• python -m py_compile autoparallel/context_parallel.py autoparallel/tracing.py autoparallel/api.py autoparallel/shardings/placement_options.py tests/test_context_parallel.py tests/test_correctness.py
• python -m pytest tests/test_context_parallel.py tests/test_correctness.py::test_correctness_attention tests/test_optimize_placement.py::test_local_map_placement_respected tests/test_activation_checkpointing.py::test_local_map_custom_metadata_propagation -q