[sraman-rgb]

mxfp8_transpose_swizzle permutes the MXFP8 scales of the attention inputs from SBHD/BSHD to BHSD with multi_tensor_transpose_to_bhsd and then swizzles them for the GEMM with multi_tensor_swizzle_scales_for_gemm_, two kernels per direction. A Triton kernel per direction now reads the source layout and writes the swizzled BHSD layout directly, for up to three tensors in one launch, with byte-identical scales.

The fused kernels are used when Triton is available, there are at most three tensors, the sequence length is a multiple of 128 and the columnwise scale rows are a multiple of 4; otherwise the existing kernels run.

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Change A
• Change B

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge; no actionable issues were found.

Summary

This PR combines the BHSD transpose and GEMM swizzle of MXFP8 attention scales to reduce kernel launches.

• MXFP8 attention scales reach GEMM-ready layout with fewer kernel launches.

No actionable issues were found. Tests were inspected but not run.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[MXFP8 attention scales] --> B{Triton available and shapes supported?}
    B -->|Yes| C[Fused BHSD transpose and swizzle]
    B -->|No| D[Existing BHSD transpose]
    D --> E[Existing GEMM swizzle]
    C --> F[Swizzled scales for attention]
    E --> F

Loading

_{Reviews (1) · Last reviewed commit: "[PyTorch] Fuse BHSD permute and GEMM swi..." · Reviewed by Greptile}