[AutoSP] Honor seq_dim when AutoSP shards inputs

[jinyouzhi]

Split from #8457.

Problem

prepare_autosp_inputs() takes a seq_dim argument, but only passes it to mark_dynamic. The AutoSP graph passes ignore it: pass_shard_seq_dim and shard_tensor_node always read the sequence length from shape[1] and slice dimension 1.

Sequence-first [seq, batch] inputs (seq_dim=0) are therefore mishandled:

• If the batch dimension is static, compilation fails with expected sequence dimension to be of type SymInt but found int.
• If the batch dimension is dynamic, AutoSP shards the batch instead of the sequence, and attention runs over the wrong tokens.

Fix

Dynamo copies only the tag attribute of an input tensor into graph metadata (_extract_tensor_dict), so seq_dim has to travel in the tag:

• prepare_autosp_inputs() tags inputs as (key, seq_dim).
• find_node_by_tag() accepts both the tuple and a bare key. A bare key, as set manually in existing code and tests, keeps the previous default of dimension 1.
• New get_autosp_seq_dim() reads the recorded dimension. pass_shard_seq_dim and shard_tensor_node use it to find the sequence length, and shard_tensor_node builds its slice for the tensor's full rank.
• label_id is now bounds-checked against seq_dim, like the other prepared inputs, since it is also sharded along that dimension.
• Docs: explain that seq_dim applies to every prepared tensor, and that seq_dim=0 is for sequence-first inputs.

Tests

TestAutoSPSeqDim (2 ranks, sp_size=2) runs a small causal SDPA model with the sequence on seq_dim ∈ {0, 1}. For each of three batches, the mean of the per-rank AutoSP losses must match a single-device run on the full sequence (rtol=1e-4). Causal attention makes the loss depend on which dimension is sharded, so slicing the batch instead of the sequence would not match.

• seq_dim=0 fails on master and passes with this PR.
• seq_dim=1 passes on both, which covers the default path with the new tuple tags.

The test compares forward losses only. AutoSP currently sums gradients over the SP group, while each rank's loss is a local mean, so gradients are sp_size× a full-sequence reference. That is part of the causal-loss aggregation follow-up from #8457, not this change.

Results on 4× NVIDIA RTX 5090 D, torch 2.13.0, transformers 5.17.0:

• tests/unit/v1/compile/test_compile_autosp.py -k "not TestAutoSPCompile": 38 passed.
• TestAutoSPCompile fails identically on unmodified master in this environment. The error is a Dynamo InternalTorchDynamoError in the Ulysses baseline's torch.compile forward, before AutoSP runs, so I excluded it.
• pre-commit passes on the changed files.

[jinyouzhi]

Could you take a look at this PR when you have a chance? Thanks @pengdurice

[tohtana]

Thank you, @jinyouzhi! Looks good to me.