NVIDIA / NVIDIA/cutlass

GemmUniversalWithBroadcast ColumnMajor-output transpose leaves the bias vector indexed in the wrong coordinate system

Open
#3,540 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

CUTLASS C++
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:

  1. a bias vector of length n (the documented per-column broadcast for the user's m x n output) is read at indices up to m - 1: out-of-bounds reads whenever m > n, and wrong elements whenever m != n;
  2. 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

  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 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.