NHD layout

[sunjiweiswift]

NHD: the last 3 dimensions are organized as (seq_len, num_heads, head_dim).
HND: the last 3 dimensions are organized as (num_heads, seq_len, head_dim).

In VLLM/sglang, NHD is a more commonly used format. Support for NHD has been added in the release pr.

[sunjiweiswift]

@petercad pls review

[sunjiweiswift]

@jiyang1011 @taozha2 @tdeng5 pls review

[Copilot]

Pull Request Overview

This PR adds support for NHD (seq_len, num_heads, head_dim) layout in addition to the existing HND (num_heads, seq_len, head_dim) layout for the BMG flash attention example. The NHD layout is commonly used in VLLM/sglang frameworks and is set as the new default.

Key changes:

• Added --layout command-line option with validation for "NHD" and "HND" values
• Updated stride calculations to support both layout formats
• Modified the verification function to handle layout-specific tensor indexing and data reordering

[sunjiweiswift]

@copilot open a new pull request to apply changes based on the comments in this thread

[sunjiweiswift]

@rolandschulz pls review and merge

[hshen14]

@sunjiweiswift did you observe some perf gains using HND than NHD when low precision e.g., FP8 is enabled? FlashInfer says it's more friendly for GPU implementation.

[sunjiweiswift]

@sunjiweiswift did you observe some perf gains using HND than NHD when low precision e.g., FP8 is enabled? FlashInfer says it's more friendly for GPU implementation.

@hshen14 There is not much difference in the case of BF16. FP8 and other low-precision types are not currently supported.