[Bug]: TND FlashAttention softmax output DTensor layout rank mismatch

Checklist

• I have searched the existing issues (https://gitcode.com/mindspore/hyper-parallel/issues).
• I have read the relevant documentation.
• I have created a reproducible case that demonstrates the issue, including the error message and traceback.

🐛 Describe the bug

Summary

With Ulysses context parallelism and FlashAttention input_layout="TND", Hyper Parallel attaches a 4-D tensor_map to the softmax_max/softmax_sum outputs, while the MindSpore kernel returns 3-D local tensors. The resulting DTensor has inconsistent metadata:

local tensor shape rank: 3
layout.tensor_map rank:  4

This becomes a runtime failure on the latest r1.0.0 branch (b36bf6006bd224a6b720f9a10dd73086b8b7b92a).

Reproduction

Environment and parallel configuration:

• Hyper Parallel: r1.0.0@b36bf6006bd224a6b720f9a10dd73086b8b7b92a
• MindSpore: 2.10.0
• MindFormers: 138f97c02b60f7ad5e759e900224d934b1601d01
• Ascend, 8 ranks
• TP=1, PP=1, CP=4
• context_parallel_method: ulysses
• FlashAttention input_layout="TND"
• max-attention-logit monitoring enabled, which queries the DTensor layout of softmax_max

The first forward step fails when converting the local softmax shape back to its global shape:

  File "mindformers/pynative/layers/flash_attention.py", line 56, in _local_head_slice
    full_shape = tuple(int(d) for d in layout.get_global_shape(local_shape))
  File "hyper_parallel/core/dtensor/layout.py", line 687, in get_global_shape
    return self._mesh.get_global_shape(slice_shape, self._tensor_map)
  File "hyper_parallel/core/dtensor/device_mesh.py", line 1291, in get_global_shape
    raise ValueError(
ValueError: Length of slice_shape (3) must match the length of tensor_map (4).

The same training case runs for multiple steps with the earlier Hyper Parallel commit 65d7e501, but fails at step 0 with b36bf600.

Root cause and regression history

MindSpore defines the FlashAttention outputs as follows:

• Non-TND: softmax_max/softmax_sum shape is (B, N, S, 8).
• TND: softmax_max/softmax_sum shape is (T, N, 8).

Therefore, the TND softmax outputs are 3-D.

However, _build_softmax_tensor_map() currently initializes a fixed four-element map and also returns four elements for TND:

softmax_tm = [-1, -1, -1, -1]
if input_layout_str == "TND":
    softmax_tm[0] = query_tm[0]
    softmax_tm[1] = query_tm[1]
return tuple(softmax_tm)

This produces a map equivalent to (T_map, N_map, -1, -1) for a real (T, N, 8) tensor.

The incorrect 4-D inference existed before the observed regression, but the old output wrapping path masked it. Previously, wrap_output() called:

DTensor.from_local(item, layout.mesh, layout.alias_placements)

DTensor.from_local() rebuilt the layout using len(local_tensor.shape), so a 3-D kernel output was wrapped with a newly constructed 3-D layout.

Commit ec7cdac9a5b17de3a56e5244080b7578f384b2bf (fix: align dtensor dispatch layout handling) changed output wrapping to reuse the inferred Layout directly:

DTensor(item, layout.mesh, layout.placements, layout)

This optimization is reasonable, but it exposes the pre-existing invalid TND softmax layout instead of implicitly rebuilding it. The resulting 3-D local tensor plus 4-D layout causes the failure above.

Expected behavior

For TND input, softmax_max and softmax_sum should receive a 3-D tensor map matching (T, N, 8):

if input_layout_str == "TND":
    return (
        query_tm[0] if len(query_tm) > 0 else -1,
        query_tm[1] if len(query_tm) > 1 else -1,
        -1,
    )

The Ulysses head sharding carried by query_tm[1] should remain intact, and len(output.layout.tensor_map) should equal len(output.to_local().shape).

Additional context

Suggested fix and validation:

1. Make _build_softmax_tensor_map() return three elements for TND softmax outputs and four elements for non-TND layouts.
2. Update the existing TND unit tests, which currently assert a 4-D softmax tensor map.
3. Add a test that wraps an actual or representative 3-D TND softmax output through wrap_output(), then verifies that local rank and tensor-map rank match.
4. Add an 8-rank Ulysses smoke test that executes at least one FlashAttention forward/backward step.

Relevant commits:

• Latest affected r1.0.0: https://gitcode.com/mindspore/hyper-parallel/commit/b36bf6006bd224a6b720f9a10dd73086b8b7b92a
• Dispatch behavior change: https://gitcode.com/mindspore/hyper-parallel/commit/ec7cdac9a5b17de3a56e5244080b7578f384b2bf

Environment info

Hardware: Ascend, 8 devices
MindSpore: 2.10.0
Hyper Parallel: r1.0.0@b36bf6006bd224a6b720f9a10dd73086b8b7b92a
MindFormers: 138f97c02b60f7ad5e759e900224d934b1601d01
Parallelism: TP1 / PP1 / CP4 Ulysses
FlashAttention layout: TND

schema_version: 1
source: gitcode
gitcode_repo: mindspore/hyper-parallel
gitcode_issue: 292
source_url: https://gitcode.com/mindspore/hyper-parallel/issues/292

Start at _build_softmax_tensor_map() and the output wrapping path, then inspect hyper_parallel/core/dtensor/layout.py and device_mesh.py for rank validation. Update the TND assertions in the existing unit tests and add coverage through wrap_output() for a 3-D softmax result. Done means the local tensor rank and tensor-map rank match, with the proposed Ulysses smoke test also passing if the environment is available.