Size SM100/SM103 array GEMM tensormap workspace from the device SM count - #3584
Open
elbourne12345 wants to merge 1 commit into
Open
Size SM100/SM103 array GEMM tensormap workspace from the device SM count#3584elbourne12345 wants to merge 1 commit into
elbourne12345 wants to merge 1 commit into
Conversation
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #3585.
Summary
The Blackwell ptr-array and grouped GEMM kernels (
sm100_gemm_array_tma_warpspecialized.hppand 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 (%smidfor ptr-array kernels, linear CTA id for grouped kernels). The workspace was sized from the rawargs.hw_info.sm_countinget_workspace_size(),initialize_workspace()andto_underlying_arguments(), and the rawhw_infowas stored inParams.to_underlying_arguments()did query the device SM count when the field was left at its default of0, but the result was only used in a trace message.Consequences with a default
KernelHardwareInfo:kArray) kernels: the launch grid is the full tile grid and does not depend onsm_count, so the kernel runs with a zero-byte tensormap workspace. Each CTA writes its 128-byte descriptors totensormaps[%smid * N]through a null or undersized pointer, and the A and B slots alias because the B offset issm_count * N = 0. This is an illegal-address fault, or silent corruption of adjacent device memory when a fusion workspace exists. A user-suppliedsm_countsmaller than the device SM count has the same effect for every SM whose%smidis at or beyond that count, and nothing rejects that configuration. The existing trace text for these kernels told users not to set the field.hw_infostored inParamscollapses the launch grid to zero, sorun()fails even thoughcan_implement()accepted the arguments, and the device query into_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 defaultKernelHardwareInfois safe on Hopper and unsafe on Blackwell.Changes
get_workspace_hw_info()helper to the four kernels. It fills in the device SM count whensm_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%smiddoes not honor a smaller user-supplied value.hw_infoconsistently inget_workspace_size(),initialize_workspace()andto_underlying_arguments(), pass it to the mainloop and tile scheduler, and store it inParamsso that the device-side descriptor stride matches the host-side sizing.Behaviour for callers that already set
hw_info.sm_countto 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
-fsyntax-only) viacutlass/gemm/kernel/gemm_universal.hpp.sm100_*ptr_array*andsm100_*group*device tests intest/unit/gemm/deviceexercise the modified functions withsm_countset, and running one of them with a defaultKernelHardwareInfowould exercise the fixed path.