Skip to content

Support zero-sized Hopper MX scale layouts#10275

Merged
roman-openai merged 1 commit into
mainfrom
roman/hopper-zero-sized-mx-layouts
May 11, 2026
Merged

Support zero-sized Hopper MX scale layouts#10275
roman-openai merged 1 commit into
mainfrom
roman/hopper-zero-sized-mx-layouts

Conversation

@roman-openai
Copy link
Copy Markdown
Collaborator

@roman-openai roman-openai commented May 9, 2026

Fix Hopper MX scale layout conversion for zero-sized scale tensors.

torch.nn.functional.pad rejects some valid zero-numel outputs. This shows up in Hopper scale swizzle after RHS transpose: for example, (0, 64) becomes (64, 0), and row padding still leaves a zero-width output.

The fix mirrors Blackwell scale layout: skip pad when there are no elements, update M/K to the padded extents, and let the existing zero-element reshape produce the swizzled storage shape.

Also relax only StridedLayoutTransformation.swizzle_data for zero-numel input. Hopper scale roundtrip can produce a valid empty canonical tensor such as (2, 0) with stride(-1) == 2; there are no elements for the packed-stride invariant to constrain.

Validation:

  • H100: PYTHONPATH=$PWD/python/triton_kernels python -m pytest -q python/triton_kernels/tests/test_tensor_details/test_layout_hopper.py::test_mxfp4_scale_zero_sized_roundtrip

@roman-openai roman-openai force-pushed the roman/hopper-zero-sized-mx-layouts branch 4 times, most recently from e332c8a to 6453855 Compare May 9, 2026 16:37
@roman-openai roman-openai force-pushed the roman/hopper-zero-sized-mx-layouts branch from 6453855 to b8df62b Compare May 9, 2026 16:51
@roman-openai roman-openai marked this pull request as ready for review May 9, 2026 17:35
@roman-openai roman-openai requested a review from ptillet as a code owner May 9, 2026 17:35
@roman-openai roman-openai merged commit 3406408 into main May 11, 2026
9 checks passed
@roman-openai roman-openai deleted the roman/hopper-zero-sized-mx-layouts branch May 11, 2026 16:03
@roman-openai roman-openai restored the roman/hopper-zero-sized-mx-layouts branch May 14, 2026 22:00
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants