GemmUniversalWithBroadcast ColumnMajor-output transpose leaves the bias vector indexed in the wrong coordinate system
Nobody has claimed this yet.
- Dominant language
- C++
- Stars
- 10.5k
- Forks
- 2.1k
- Avg merge
- 3d 11h
- Merged PRs (30d)
- 7
Description
Description
The ColumnMajor-output specializations of device::GemmUniversalWithBroadcast and device::GemmUniversalStreamKWithBroadcast implement their output layout by transposing the problem (to_underlying_arguments calls args.transposed_problem()), but the transpose does not touch the broadcast-vector arguments even though they are consumed in the transposed coordinate system.
// include/cutlass/gemm/kernel/gemm_with_fused_epilogue.h:208-217
Arguments transposed_problem() const {
Arguments args(*this);
std::swap(args.problem_size.m(), args.problem_size.n());
std::swap(args.ptr_A, args.ptr_B);
std::swap(args.lda, args.ldb);
std::swap(args.batch_stride_A, args.batch_stride_B);
return args; // ptr_Vector / ldr / batch_stride_Vector untouched
}
while the fused epilogue indexes the bias by the underlying kernel's output columns:
// include/cutlass/gemm/kernel/gemm_with_fused_epilogue.h:630-631 (and :762-763)
if (ptr_Vector) {
ptr_Vector += threadblock_offset.column() + threadblock_tile_offset.m() * params.ldr;
}
After the transpose, the underlying kernel's column extent is the user's M, so:
- a bias vector of length
n(the documented per-column broadcast for the user'sm x noutput) is read at indices up tom - 1: out-of-bounds reads wheneverm > n, and wrong elements wheneverm != n; ldr, which scales the per-M-tile offset, keeps its pre-transpose meaning relative to an axis that no longer corresponds.
gemm_universal_with_broadcast.h:315-317 and gemm_universal_streamk_with_broadcast.h:315-317 both route through this transpose. Nothing in-tree exercises the ColumnMajor-output specialization with ptr_Vector set, which is why this has gone unnoticed.
Suggested fix
Decide and document the post-transpose convention for the vector: either swap the vector pointer semantics alongside the problem (and adjust the epilogue indexing), or reject ptr_Vector for the ColumnMajor-output specializations instead of silently misindexing.
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with include/cutlass/gemm/kernel/gemm_with_fused_epilogue.h, especially Arguments::transposed_problem() and the epilogue indexing at lines 630-631 and 762-763. Trace the ColumnMajor-output paths in gemm_universal_with_broadcast.h and gemm_universal_streamk_with_broadcast.h, then add coverage for ptr_Vector with unequal dimensions. Done means the vector follows a documented post-transpose convention or is rejected safely.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- hpc, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100