SMEM offset caching RHT

[sraman-rgb]

Description

Cache offsets and first_dims in shared memory for graph-safe kernel, helps with reduction of memory bytes read

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]

Greptile Summary

This PR caches offsets and first_dims from global memory into shared memory via cp.async (LDGSTS) in the DMA warp, then uses mbarrier synchronization to ensure the data is visible before epilogue warps read it. The barrier strategy is well-designed: for kEnableRHTColQuant=true, tma_barrier[0] is promoted from 1 to 2 required arrivals (DMA cp.async completion + TMA B hardware) so both data-readiness conditions share a single wait; for kEnableRHTColQuant=false, a new cpasync_barrier[0] carries just the cp.async signal to the row-quant warp without interfering with TMA B synchronization.

The memory ordering is correct — cp_async_wait<0>() + __threadfence_block() + mbarrier_arrive (release) paired with mbarrier_wait (acquire) in consumers is the standard LDGSTS→mbarrier pattern. One remaining gmem read that the smem cache could eliminate is noted inline.

Confidence Score: 5/5

Safe to merge; the barrier logic and memory ordering are correct and the only finding is a P2 missed-optimization suggestion.

All barrier initialization, arrive, and wait sequences are correct for both the kEnableRHTColQuant=true (2-arrival tma_barrier) and false (cpasync_barrier) paths. The __syncthreads after barrier init, the cp_async_wait+threadfence_block release pattern, and the constexpr-guarded wait selection are all sound. The sole inline comment is a P2 optimization opportunity (one per-tile gmem read of offsets[group_idx] that could use offsets_smem), not a correctness defect.

graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu — verify the per-tile offsets[group_idx] access at line 869 is intentionally left as gmem or should be switched to offsets_smem.

Important Files Changed

Filename	Overview
transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu	Adds SMEM caching of offsets/first_dims via cp.async with correct barrier synchronization; one per-tile gmem read of offsets[group_idx] at line 869 remains in the hot path despite the smem cache being available.
transformer_engine/common/hadamard_transform/test_offset_caching.cu	New GTest unit test verifying cp.async-cached offsets match gmem results; build system integration noted as pre-existing concern.
tests/pytorch/nvfp4/bench_graph_safe_swizzle.py	New benchmark script for graph-safe swizzle; module-level GPU execution noted as pre-existing concern.
tests/pytorch/nvfp4/bench_structural.py	Structural benchmark covering O(1), binary-search, and linear-scan paths; module-level GPU execution noted as pre-existing concern.
tests/pytorch/nvfp4/bench_search.py	New benchmark comparing graph-safe equal/unequal splits and non-graph-safe quantization; module-level GPU execution noted as pre-existing concern.
tests/pytorch/nvfp4/bench_sweep_swizzle.py	Sweeps swizzle ON/OFF across M values; module-level GPU execution noted as pre-existing concern.
tests/pytorch/nvfp4/ncu_test.py	Minimal NCU profiling harness with warmup + single measured launch; module-level GPU execution noted as pre-existing concern.

Sequence Diagram

sequenceDiagram
    participant DMA as DMA Warp
    participant ColQ as Col Quant Warp
    participant RowQ as Row Quant Warp
    participant TMAB as TMA B Hardware

    Note over DMA: cp.async offsets[] → smem_offsets<br/>cp.async first_dims[] → smem_first_dims
    DMA->>DMA: cp_async_fence() + cp_async_wait<0>() + __threadfence_block()

    alt kEnableRHTColQuant=true
        DMA->>DMA: mbarrier_arrive(tma_barrier[0])  ①
        DMA->>TMAB: launch TMA B load
        TMAB-->>DMA: (async)
        TMAB->>ColQ: mbarrier_arrive(tma_barrier[0])  ②
        Note over ColQ: wait_barrier(tma_barrier[0], phase=0)<br/>fires after both ① and ②
        Note over RowQ: wait_barrier(tma_barrier[0], phase=0)<br/>fires after both ① and ②
        ColQ->>ColQ: read offsets_smem / first_dims_smem
        RowQ->>RowQ: read offsets_smem / first_dims_smem
    else kEnableRHTColQuant=false
        DMA->>DMA: mbarrier_arrive(cpasync_barrier[0])  ①
        DMA->>TMAB: launch TMA B load
        TMAB->>TMAB: mbarrier_arrive(tma_barrier[0]) - no waiter
        Note over RowQ: wait_barrier(cpasync_barrier[0], phase=0)<br/>fires after ①
        RowQ->>RowQ: read offsets_smem / first_dims_smem
    end

Loading

_{Reviews (5): Last reviewed commit: "Add copyright headers to nvfp4 benchmark..." | Re-trigger Greptile}

[greptile-apps]

Module-level GPU code will execute on pytest import

All five new scripts (bench_structural.py, bench_sweep_swizzle.py, bench_search.py, bench_graph_safe_swizzle.py, ncu_test.py) contain GPU kernel launches at module scope. When pytest discovers files in tests/pytorch/nvfp4/, it imports each one to collect tests; the imports execute the benchmarks immediately — potentially hanging or crashing CI on machines without the required GPU or package.

Wrap the benchmark body in a if __name__ == "__main__": guard on all five files, e.g.:

if __name__ == "__main__":
    for M in [16384, 65536, 131072]:
        ...

[CalebDu]

use NumThreadsPerWarp in cutlass header. Do not need redundant definition.

[CalebDu]

use cp_async_fence() and cp_async_wait().

[greptile-apps]

Test file never runs — missing build system integration

test_offset_caching.cu uses GTest (TEST(...) fixtures) but has no main() function and is not referenced in any CMakeLists.txt (there is no CMakeLists.txt in transformer_engine/common/hadamard_transform/, and a project-wide search finds no reference to this file). As a result this test is never compiled or executed in CI, so the barrier correctness it is meant to verify is untested.

To wire it in, either add a CMakeLists.txt entry that links against gtest_main, or add a main() directly:

int main(int argc, char **argv) {
  ::testing::InitGoogleTest(&argc, argv);
  return RUN_ALL_TESTS();
}