NVIDIA / NVIDIA/cutlass

[QST] How to build a pipeline between epilogue and communication warp groups?

Open
#2,900 3 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

? - Needs Triage inactive-30d inactive-90d question
Dominant language
C++
Stars
10.5k
Forks
2.1k
Avg merge
3d 11h
Merged PRs (30d)
7

Description

Hi, recently I am working on the examples of CuTe DSL.

In hopper/dense_gemm_persistent.py, the epilogue is like as the following:

                num_prev_epi_tiles = tile_sched.num_tiles_executed * epi_tile_num
                for epi_idx in cutlass.range_constexpr(epi_tile_num):
                    # Copy from accumulators to D registers
                    for epi_v in cutlass.range_constexpr(size_tRS_rD):
                        tRS_rD[epi_v] = tRS_rAcc[epi_idx * size_tRS_rD + epi_v]

                    # Type conversion
                    acc_vec = tRS_rD.load()
                    tRS_rD_out.store(acc_vec.to(self.c_dtype))

                    # Copy from D registers to shared memory
                    epi_buffer = (num_prev_epi_tiles + epi_idx) % cute.size(
                        tRS_sD, mode=[3]
                    )
                    cute.copy(
                        tiled_copy_r2s,
                        tRS_rD_out,
                        tRS_sD[(None, None, None, epi_buffer)],
                    )

                    cute.arch.fence_proxy(
                        cute.arch.ProxyKind.async_shared,
                        space=cute.arch.SharedSpace.shared_cta,
                    )
                    self.epilog_sync_barrier.arrive_and_wait()

                    gmem_coord = epi_tile_layout.get_hier_coord(epi_idx)
                    
                    # Copy from shared memory to global memory
                    if warp_idx == self.epi_store_warp_id:
                        cute.copy(
                            tma_atom_c,
                            bSG_sD[(None, epi_buffer)],
                            bSG_gD[(None, gmem_coord)],
                        )
                        tma_store_pipeline.producer_commit()
                        tma_store_pipeline.producer_acquire()

                    self.epilog_sync_barrier.arrive_and_wait()

Since it uses a multi-stage epilogue pipeline, when calling tma_store_pipeline.producer_acquire(), it actually calls cute.arch.cp_async_bulk_wait_group(num_stages-1) under the hood. However, this becomes complex when I need to use another specialized warp to consume the results finished by MMA warp groups, because I cannot know whether the epilogue pipeline is initially empty or a previous tile has finished.

In addition, tma_store_pipeline.producer_tail() can indeed ensure all in-flight TMA Store transactions are finished, but it's called after all tiles are done, which means the communication warp cannot work parallel with the MMA warp groups.

Essentially, what I want to implement is something like this:

### Producer(Epilogue) ###
while work_tile.is_valid_tile:
    # Wait until the flag is empty
    producer.acquire()

    ## Epilogue, TMA store to GMEM
    cute.copy(tma_atom_c, bSG_sD, bSG_gD)
   
    # Signal the consumer tile is ready.
    producer.commit()

### Consumer ###
while work_tile.is_valid_tile:
    # Wait until the tile is ready.
    consumer.wait()

    ## Do something with the tile, such as communications

    # Release the tile
    consumer.release()

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with hopper/dense_gemm_persistent.py and trace tma_store_pipeline producer_acquire, producer_commit, producer_tail, and the epilogue synchronization barriers. Determine whether the requested producer-consumer protocol is supported for communication warp groups, and validate any proposed behavior with a focused CuTe DSL example or test showing overlap and correct tile handoff.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python
Domain
hpc, performance
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.