Skip to content

Size SM100/SM103 array GEMM tensormap workspace from the device SM count - #3584

Open
elbourne12345 wants to merge 1 commit into
NVIDIA:mainfrom
elbourne12345:fix-sm100-array-tensormap-workspace-sm-count
Open

Size SM100/SM103 array GEMM tensormap workspace from the device SM count#3584
elbourne12345 wants to merge 1 commit into
NVIDIA:mainfrom
elbourne12345:fix-sm100-array-tensormap-workspace-sm-count

Conversation

@elbourne12345

@elbourne12345 elbourne12345 commented Sep 4, 2026

Copy link
Copy Markdown

Fixes #3585.

Summary

The Blackwell ptr-array and grouped GEMM kernels (sm100_gemm_array_tma_warpspecialized.hpp and its input-transform, mma-transform and SM103 block-scaled siblings) keep one set of TMA descriptors per SM in the global workspace and index it on the device by SM id (%smid for ptr-array kernels, linear CTA id for grouped kernels). The workspace was sized from the raw args.hw_info.sm_count in get_workspace_size(), initialize_workspace() and to_underlying_arguments(), and the raw hw_info was stored in Params. to_underlying_arguments() did query the device SM count when the field was left at its default of 0, but the result was only used in a trace message.

Consequences with a default KernelHardwareInfo:

  • ptr-array (kArray) kernels: the launch grid is the full tile grid and does not depend on sm_count, so the kernel runs with a zero-byte tensormap workspace. Each CTA writes its 128-byte descriptors to tensormaps[%smid * N] through a null or undersized pointer, and the A and B slots alias because the B offset is sm_count * N = 0. This is an illegal-address fault, or silent corruption of adjacent device memory when a fusion workspace exists. A user-supplied sm_count smaller than the device SM count has the same effect for every SM whose %smid is at or beyond that count, and nothing rejects that configuration. The existing trace text for these kernels told users not to set the field.
  • grouped kernels: the raw hw_info stored in Params collapses the launch grid to zero, so run() fails even though can_implement() accepted the arguments, and the device query in to_underlying_arguments() is dead code.

The SM90 array kernels (sm90_gemm_array_tma_warpspecialized_cooperative.hpp, _pingpong.hpp) already query the device SM count and size the workspace from it, so a default KernelHardwareInfo is safe on Hopper and unsafe on Blackwell.

Changes

  • Add a get_workspace_hw_info() helper to the four kernels. It fills in the device SM count when sm_count <= 0, and for kernels that index the workspace by %smid (the non-grouped SM100 path and the input-transform kernel) it never sizes the workspace below the device SM count, since %smid does not honor a smaller user-supplied value.
  • Use the resulting hw_info consistently in get_workspace_size(), initialize_workspace() and to_underlying_arguments(), pass it to the mainloop and tile scheduler, and store it in Params so that the device-side descriptor stride matches the host-side sizing.
  • Reword the trace messages that described the field as a performance hint.

Behaviour for callers that already set hw_info.sm_count to the device SM count (all in-tree tests, examples and the profiler) is unchanged: the helper returns the same value and the same workspace size.

Testing

  • The four headers parse cleanly with GCC 15 in a host-only build (-fsyntax-only) via cutlass/gemm/kernel/gemm_universal.hpp.
  • I do not have Blackwell hardware, so I could not run the SM100/SM103 unit tests; the change is a host-side sizing fix and the device code is untouched. The existing sm100_*ptr_array* and sm100_*group* device tests in test/unit/gemm/device exercise the modified functions with sm_count set, and running one of them with a default KernelHardwareInfo would exercise the fixed path.

The Blackwell ptr-array and grouped GEMM kernels keep one set of TMA
descriptors per SM in the workspace and index it on the device by SM id
(%smid for ptr-array kernels, linear CTA id for grouped kernels). The
workspace was sized from the raw args.hw_info.sm_count in
get_workspace_size(), initialize_workspace() and to_underlying_arguments(),
and the raw hw_info was stored in Params. to_underlying_arguments() queried
the device SM count when the user left the field at its default of 0, but
the result was only traced and never used.

With a default KernelHardwareInfo the ptr-array kernels therefore got a
zero-byte tensormap workspace while still launching over the full tile
grid, so every CTA wrote its 128-byte descriptors through a null or
undersized pointer, and the A and B descriptor slots aliased because the B
offset is sm_count * NumTmaDescriptorsPerSm. A user-supplied sm_count
smaller than the device SM count had the same effect for the SMs whose
%smid is at or beyond that count. For grouped kernels the same default
collapsed the launch grid to zero, so the launch failed even though
can_implement() had accepted the arguments. The SM90 array kernels already
query the device SM count and size the workspace from it.

Add a get_workspace_hw_info() helper to the SM100 array, input-transform,
mma-transform and SM103 block-scaled array kernels that fills in the device
SM count when none was supplied and, for kernels that index by %smid,
never sizes the workspace below the device SM count. Use the resulting
hw_info consistently in all three workspace functions and store it in
Params so that the device sees the same stride, and reword the trace
messages that described the field as optional.

Signed-off-by: Aamir Ahmed <elb12345@hotmail.co.uk>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

1 participant