universal_gemm_kernel.hpp:581-584 rejects any D tensor whose layout is not CLayout — and does so silently: no log line, no assert message, the kernel simply does not run. Localizing it took a CK_TILE_LOGGING run plus a source read.
This blocks a natural and otherwise-supported composition: a weight-preshuffled GEMM plus a per-token/per-channel scale epilogue, expressed as two broadcast D tensors —
- D0 = per-token A scale: an M-vector,
ColumnMajor, stride 0 in N
- D1 = per-channel B scale: an N-vector,
RowMajor, stride 0 in M
D1 is accepted; D0 is not, because its layout differs from CLayout. Independently, cshuffle_epilogue.hpp:782 wraps every D window with the (M, N) distribution and has no handling for a column-layout D.
This is exactly the architecture CK's own legacy XDL path ships (GridwiseGemmMultiD with row/col scale Ds), so the operation is well-precedented; only the CK Tile universal path can't express it.
Asks, in increasing order of effort:
- make the rejection loud (a static_assert message or a log line naming the offending D index and layout);
- relax the check and add column-layout D distribution handling in the CShuffle epilogue;
- or, if RowColQuant is the intended route for this shape of problem, document that — it works (see the
QuantGemmKernel RowColQuant path) and it is not obvious from the examples.
Context: found while wiring per-token/per-channel FP8 GEMMs on gfx950; the RowColQuant + weight-preshuffle route works and measured +8-11% over the flatmm path at three GEMM shapes, so the composition is worth having.
universal_gemm_kernel.hpp:581-584rejects any D tensor whose layout is notCLayout— and does so silently: no log line, no assert message, the kernel simply does not run. Localizing it took aCK_TILE_LOGGINGrun plus a source read.This blocks a natural and otherwise-supported composition: a weight-preshuffled GEMM plus a per-token/per-channel scale epilogue, expressed as two broadcast D tensors —
ColumnMajor, stride 0 in NRowMajor, stride 0 in MD1 is accepted; D0 is not, because its layout differs from CLayout. Independently,
cshuffle_epilogue.hpp:782wraps every D window with the (M, N) distribution and has no handling for a column-layout D.This is exactly the architecture CK's own legacy XDL path ships (
GridwiseGemmMultiDwith row/col scale Ds), so the operation is well-precedented; only the CK Tile universal path can't express it.Asks, in increasing order of effort:
QuantGemmKernelRowColQuant path) and it is not obvious from the examples.Context: found while wiring per-token/per-channel FP8 GEMMs on gfx950; the RowColQuant + weight-preshuffle route works and measured +8-11% over the flatmm path at three GEMM shapes, so the composition is worth having.