Skip to content

[tritonbench][AutoWS] Enable AutoWS matmul kernels on Hopper#1025

Open
njriasan wants to merge 1 commit intometa-pytorch:mainfrom
njriasan:export-D101840063
Open

[tritonbench][AutoWS] Enable AutoWS matmul kernels on Hopper#1025
njriasan wants to merge 1 commit intometa-pytorch:mainfrom
njriasan:export-D101840063

Conversation

@njriasan
Copy link
Copy Markdown
Contributor

Summary:
Extends the Blackwell-only TMA / persistent / descriptor matmul benchmarks
in tritonbench to also run on Hopper.

Changes:

  1. Widen enabled=IS_BLACKWELL to enabled=IS_BLACKWELL or IS_HOPPER for
    the 6 triton_blackwell_* benchmarks in operator.py. The underlying
    kernels (matmul_kernel_tma, matmul_kernel_tma_persistent, and
    matmul_kernel_descriptor_persistent) are arch-agnostic, so they run on
    Hopper once the gates allow it.

  2. Update _prune_tma_persistent_configs in warp_spec_persistent_matmul.py
    so the FLATTEN requirement is computed correctly across {Hopper,
    Blackwell} x {meta WS, OAI Triton WS, no WS}. tl.range(flatten=True) is
    only supported on Blackwell with the OAI Triton WS path; on Hopper we now
    require FLATTEN=False for both AutoWS (meta WS) and OAI Triton paths.

Differential Revision: D101840063

Summary:
Extends the Blackwell-only TMA / persistent / descriptor matmul benchmarks
in `tritonbench` to also run on Hopper.

Changes:
1. Widen `enabled=IS_BLACKWELL` to `enabled=IS_BLACKWELL or IS_HOPPER` for
   the 6 `triton_blackwell_*` benchmarks in `operator.py`. The underlying
   kernels (`matmul_kernel_tma`, `matmul_kernel_tma_persistent`, and
   `matmul_kernel_descriptor_persistent`) are arch-agnostic, so they run on
   Hopper once the gates allow it.

2. Update `_prune_tma_persistent_configs` in `warp_spec_persistent_matmul.py`
   so the `FLATTEN` requirement is computed correctly across {Hopper,
   Blackwell} x {meta WS, OAI Triton WS, no WS}. `tl.range(flatten=True)` is
   only supported on Blackwell with the OAI Triton WS path; on Hopper we now
   require `FLATTEN=False` for both AutoWS (meta WS) and OAI Triton paths.

Differential Revision: D101840063
@meta-codesync
Copy link
Copy Markdown

meta-codesync Bot commented Apr 21, 2026

@njriasan has exported this pull request. If you are a Meta employee, you can view the originating Diff in D101840063.

@njriasan njriasan changed the title Enable Blackwell matmul tests on Hopper [tritonbench][AutoWS] Enable AutoWS matmul kernels on Hopper Apr 21, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant