Repository navigation
fix(hipblaslt): compact MX scale stride and exact rocRoller kernel-name matching - #14043
Conversation
github-actions
Bot
added
project: hipblaslt
project: hipsparselt
ci:hipsparselt-fast
labels
✅ All Checks Passed — Ready for Review
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
|
🎉 All checks passed! This PR is ready for review. |
jaopaulolc
changed the title
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## develop #14043 +/- ##
===========================================
+ Coverage 77.42% 77.46% +0.04%
===========================================
Files 2903 2906 +3
Lines 490638 490550 -88
Branches 74225 74253 +28
===========================================
+ Hits 379863 379993 +130
+ Misses 85771 85539 -232
- Partials 25004 25018 +14
*This pull request uses carry forward flags. Click here to find out more.
🚀 New features to boost your workflow:
|
a7cedfa to
c90b64d
Compare
c90b64d to
052114a
Compare
newling
left a comment
newling
left a comment
There was a problem hiding this comment.
This PR resolves an issue I was hitting; see closed #14370, which this supersedes. Some agent review comments here: https://github.com/newling/agent-public-reviews/blob/82b2dd62aca2c288c7ba8be7ca02ff9f7c274f9d/rocm-libraries/review_pr14043.md, with suggested changes linked. Feel free to take any of those, they’re nonblocking from my perspective. Approving in advance.
052114a to
f9a41a5
Compare
|
Took the two nonblocking suggestions from the review of
The description already matches this head: both the row stride and the batch stride follow |
Alex-Vasile
left a comment
Alex-Vasile
left a comment
There was a problem hiding this comment.
Looks good. One minor thing.
Two independent library correctness fixes for externally-captured custom kernels. generateCustomCall passed StrideScaleA0/StrideScaleB0 the K-block dimension that setMXScaleA/B pads with RoundUpToMultiple(CeilDivide(K, mxBlock), 8); rocRoller builds its kernels against the compact K/32, so the two agree only when K % 256 == 0. Pass the compact stride for externally-captured custom kernels only -- Tensile's own generated gfx950 MX kernels overwrite this pitch from SizesSum in their prologue, so the padded kernarg is dead for them and must stay. matchesRocRollerSelectedKernel tested the bare substring "WGT_<m>x<n>x<k>_UR_", which occurs in both "_WGT_..._UR_" and "_SK_WGT_..._UR_", so a plain tile matched a StreamK kernel and the reverse; and the prefix match stopped at "_Float_" while the tile match started at "WGT_", leaving the scale tokens between them uninspected, so a pre-swizzled kernel matched a non-pre-swizzled problem. Anchor both ends and compare the full name. Also apply the workspace reject that the rocRoller-on path already performs.
generateCustomCall now hands an external custom kernel the compact ceil(K/mxBlock) scale pitch unconditionally, but only hipBLASLt packs that buffer. The TensileLite benchmark client allocates the MX scale tensor straight from the padded descriptor and restrides the canonical scales into it, zeroing the padding tail, then hands that same buffer to the custom kernel. The compact pitch sent every odd scale row into the zeroed tail -- an E8M0 0x00 is 2^-127 -- so every MX custom-kernel config in tensilelite/Tests/custom (custom_rr_mx, custom_rr_f6, custom_rr_wgm, custom_rr_mx_pure) failed validation on gfx950 with ~3/4 of the output tensor off by 2^-127. Record the layout on the problem instead of assuming it: setCustomKernelMXScalesArePacked, set by tensile_host.cpp because the user's plain Block_32_UE8M0 pointer really is packed, left false everywhere else so the descriptor's own stride is used.
setCustomKernelMXScalesArePacked was a standalone statement next to, but separable from, the setMXScaleA/B calls it describes. Dropping it would not have failed anything: the TensileLite benchmark client leaves the flag false on purpose and keeps passing, while hipBLASLt callers would go back to receiving the padded pitch against a packed buffer -- the silent corruption the previous commit fixed. Pair the two in setHostMXScaleA/B (new include/mx_scale_setup.hpp), which every MX scale attachment in tensile_host.cpp now goes through, so there is no separate statement left to drop. clients/tests/unit/ mx_scale_setup_gtest.cpp pins the pairing: removing the declaration from either helper fails HostUnitMXScaleSetup. ConstructTensileProblem and updateTensileProblem cannot be driven from a host-only test -- both are file-static, and both call hipGetDevice through mxScaleTensorNeedsPaddingFreeDim, while these tests run with no GPU. Dropping the setup at either call site is caught at compile time instead: padMXScaleTensorFreeDim is its only consumer there, and -Wunused-variable is promoted to an error around it, the tripwire idiom clients/tests/unit already uses for -Wswitch. hipblaslt-test-hostunit names roc::tensilelite-host, which it already linked transitively, for the generated tensilelitehost/export.h that Tensile/ContractionProblem.hpp includes.
getRocRollerSelectedSolutions matched a checked-in rocRoller kernel by name alone. The name carries only the tile shape and the operand types, so a one-sided-MX problem matched a two-sided kernel and faulted the GPU, and size-window and scale-layout predicates never ran. Require hardwarePredicate and problemPredicate to accept the problem before a name match is taken, and keep searching when they do not.
StrideScaleA0/B0 is the row stride of the buffer the caller allocated. The batch stride follows that same buffer, and the comment says so.
The batch-stride test only covered the legacy argument list. Data-parallel version 1, initial strides, and a partial block use different slots.
setHostMXScaleA/B declared every hipBLASLt MX buffer packed, including a Block_32_UE8M0_32_8_EXT one. That caller pre-swizzled its scales into the padded gfx950 layout the descriptor already describes, so an external custom kernel would have been handed the compact pitch for a padded buffer. A single flag on the problem also could not describe an A and a B that arrive in different formats: whichever helper ran last decided the stride for both. Record the layout per operand. The helpers take scalesArePacked, which tensile_host.cpp derives from the operand's own format, and generateCustomCall reads A's flag for StrideScaleA0/A1 and B's for StrideScaleB0/B1. setCustomKernelMXScalesArePacked stays as the way to declare both, for a caller that lays them out alike. This is latent today: no PSTA/PTA kernel is checked in.
PR Pundit review(Testing PR Pundit Code Review agent. Not all comments may be relevant. This is a one-time automated pass: please address what is useful, no need to request a re-review. Please reply to each point, whether it applies or not: we harvest these replies, and both kinds are enormously helpful for tuning the agent. Reach out to @peymanr for suggestions.)
1. 🟠 Should fix: Predicate and workspace eligibility filters added without any test
The eligibility filters this PR adds to Who is affected: Only builds with Checked against the code by an independent pass. Related past reviews: #1944 on Shared from a pilot of PR Pundit, an agentic code reviewer that builds and tests the change and reads the project's past reviews. Reviewed at |
PR Pundit review(Testing PR Pundit Code Review agent. Not all comments may be relevant. This is a one-time automated pass: please address what is useful, no need to request a re-review. Please reply to each point, whether it applies or not: we harvest these replies, and both kinds are enormously helpful for tuning the agent. Reach out to @peymanr for suggestions.)
1. 🟠 Should fix: Format-to-layout regression in hostMXScalesArePacked passes host unit tests unnoticed
The format-to-layout mapping in Who is affected: Maintainers changing the MX format mapping later; no current user-visible configuration is affected. Checked against the code by an independent pass. Related past reviews: #4599 on Shared from a pilot of PR Pundit, an agentic code reviewer that builds and tests the change and reads the project's past reviews. Reviewed at |
TheRock Submodule Bump ActivityNewest first
|
JIRA ID : AIHPBLAS-5261
Motivation
Two independent defects on gfx950 MX return wrong numbers: the row pitch hipBLASLt reports for the scale tensor does not match the layout the kernel was compiled against, and the test that picks a kernel by inspecting its name accepts kernels that do not fit the problem. Measured over 108 gfx950 MX problem and type combinations, 36 produced numerically wrong results, some containing
inf.Technical Details
This pull request is stacked on #14086.
The first fix concerns the scale row pitch.
ContractionProblemGemm::setMXScaleA/Brounds the scale tensor's K-block dimension up withRoundUpToMultiple(CeilDivide(K, mxBlock), 8)on gfx950.generateCustomCallpassed that padded value as the kernel argumentsStrideScaleA0andStrideScaleB0, but the rocRoller kernels are built against the unpaddedceil(K/32). The two values are equal only whenK % 256 == 0. Otherwise every scale row after the first resolves to the wrong bytes, and the kernel reads past the end of the caller's scale buffer.generateCustomCallnow takes both strides from the buffer the caller passed.customKernelMXScaleAIsPacked()andcustomKernelMXScaleBIsPacked()name that caller, one operand at a time, because hipBLASLt takes the A and B scale formats independently. When an operand's flag is set, its row stride is the unpaddedceil(K/mxBlock). When it is not, the descriptor stride is already the one the kernel must use.The padding is deliberately not removed globally. The gfx950 MX kernels that TensileLite generates overwrite this pitch in their own prologue from the problem's
SizesSumfield, so the kernel argument is dead for them and the padded descriptor must stay for the rest of the library. The change is confined togenerateCustomCall, which runs only for externally captured custom kernels. The batch strideStrideScale{A,B}1follows the same caller buffer: a packed caller gets the product of the unpadded scale sizes, and an unpacked caller keeps the descriptor stride. No shipped rocRoller kernel reads that batch stride, and those kernels run a single batch. This also repairs the FP6 custom kernels added by #12964 and #12965, which reach the same code path.hipBLASLt declares each operand packed unless its format is
Block_32_UE8M0_32_8_EXT: that caller pre-swizzled its scales into the padded gfx950 layout, which the descriptor already describes. The TensileLite benchmark client leaves both flags false, because it allocates from the padded descriptor.The second fix concerns
matchesRocRollerSelectedKernel, the predicate that decides whether a kernel's name describes the problem at hand. It has three defects, all reachable on theHIPBLASLT_ENABLE_ROCROLLER=OFFbuild:WGT_<m>x<n>x<k>_UR_, which occurs inside both_WGT_..._UR_and_SK_WGT_..._UR_. A plain tile therefore matched a StreamK kernel — StreamK is a work-partitioning scheme that splits the summation axis across workgroups — and a StreamK tile matched a plain kernel. The test is now anchored as_WGT_...with an explicit_SKdiscriminator._Float_while the tile match began atWGT_, so the scale-layout tokens between the two were never inspected. A pre-swizzled kernel, named withPSTA_,PTA_,PSTB_, orPTB_, could match a problem whose scales are not pre-swizzled. Those tokens are now rejected._WGM_, becauseparameter_selection.cppforcesworkgroupMappingDim = -1for StreamK. The workgroup-mapping comparison now runs only whenworkgroupMapping && !streamK, and StreamK tiles are no longer rejected outright.getRocRollerSelectedSolutionsalso gains the workspace rejection that the rocRoller-enabled path already applies atgemm.cpp:703-706. Without it anOFFbuild hands back a StreamK kernel that anONbuild would have skipped, because_convertToHeuristicResultArrayreportsrequiredWorkspaceSizebut never filters on it.Name matching is not enough on its own
A kernel's name carries its tile shape and its operand types, and nothing else. Every other condition the kernel was captured under — which scale layout it reads, the size multiples and size windows it requires, how many batches it handles, how large a tensor it can address — lives in the kernel's selection predicates. An ordinary library lookup evaluates those; this hand-rolled search skipped them. A problem that block-scales only one of its two operands therefore matched a kernel built to read two scale tensors, and that kernel read a scale buffer the caller had never allocated.
getRocRollerSelectedSolutionsnow requires a candidate'shardwarePredicateandproblemPredicateto accept the problem before it takes a name match, and keeps searching when they do not. It tests hardware and problem only, as the grouped-gemm path in the same file does:taskPredicateneeds aTask, which needs the workspace size that is tested immediately afterwards.This narrows what the
OFFpath selects. Where the kernel it used to match declares a size window that excludes the problem, it now picks a different one. That is the intended behaviour, but it has a cost, measured below: the only candidates on this path are other rocRoller kernels, so when the best-fitting tile is correctly rejected there may be nothing comparable left in the pool.The predicate moves to
rocroller_kernel_match.hpp, whose only dependency is<string>, so it can be tested without a problem descriptor or a GPU.No longer rejecting StreamK outright changes which kernel an
OFFbuild selects. That is intended — it is what lets anOFFbuild reach the StreamK kernels at all — but this pull request tests it only at the name-matching level. The selection-identity sweep that compares anOFFbuild against anONbuild lands with the logic that makes these kernels selectable, later in this stack.Test Plan
CustomKernelTest.ScaleStrideIsCompactNotPadded,CustomKernelTest.ScaleStrideIsCompactForEitherBoundDimension,CustomKernelTest.ScaleBatchStrideFollowsTheCallerBuffer,CustomKernelTest.ScaleStrideFollowsDescriptorWhenCallerDoesNotPack,CustomKernelTest.ScaleStridesIncludeTheInitialStrideForDataParallelArgs, andCustomKernelTest.ScaleStrideRoundsUpPartialBlocks: assert thatgenerateCustomCallemits the caller's scale stride. A packed caller with aKthat is not a multiple of 256 getsceil(K/mxBlock), for both bound dimensions and for the batch stride. An unpacked caller gets the descriptor stride.HostUnitRocRollerKernelMatch.*: five host-only tests covering the StreamK and plain tile discrimination, the pre-swizzle token rejection, and the workgroup-mapping comparison. Mutation-check them by running them against the pre-patch predicate; a test that passes against the old code does not discriminate.CustomKernelTest.ScaleStrideFollowsEachOperandsBuffer: a packed A with a padded B, and the reverse, each get their own buffer's row and batch strides.HostUnitMXScaleSetup.smoke_PreSwizzledDeclarationIsPerOperand: a pre-swizzled operand is declared unpacked, and declaring one operand does not change the other.scaleA 3,scaleB 0, bf8 inputs, 1024x1024x1024 — on anHIPBLASLT_ENABLE_ROCROLLER=OFFbuild. It must report no solution rather than launching a kernel.--verifyflag on the command line does not reach a record supplied by--yaml.Test Result
Completed evidence:
CustomKernelTestcases are new and pass.HostUnitRocRollerKernelMatchtests are new and pass. Against the pre-patch predicate, 4 of the 5 fail, so they discriminate the change.infin some cases. After the change 0 are wrong and none regressed; the 63 combinations that were already clean are bit-identical.Predicate evaluation, measured on MI355X (gfx950, PCI 0x75a3) against an
HIPBLASLT_ENABLE_ROCROLLER=OFFbuild, with the numerical check set in the test record:scaleA 3andscaleB 0: before the change the run ends inan illegal memory access was encountered (700); after it, the problem reports no solution, on 2 of 2 runs.RR_GEMM_..._WGT_32x32x64_UR_2declares M and N between 64 and 128 and K exactly 128, and was being selected at 512x512x512.RR_GEMM_..._WGT_256x128x128_UR_2, whose declared window ends at M = 2560, toRR_GEMM_..._WGT_16x16x128_UR_2, and the kernel time rises by a factor of about 3.6 over three repeats. That fallback, routing a problem the rocRoller path declines to ordinary Tensile selection, is feat(hipblaslt): ship rocRoller MX kernels and packed client scales #14497, later in this stack. It is not in this pull request.Pending: automated GitHub checks have not reported on this branch. The parity sweep was run locally on gfx950 hardware and does not run in a shared CI lane today.
Submission Checklist
Risk level
Medium. Two pieces of device-visible behavior change. The scale stride passed to custom MX kernels on gfx950 changes value whenever
K % 256 != 0, which is the defect being fixed; kernels generated by TensileLite are unaffected because they overwrite that argument. Separately, anHIPBLASLT_ENABLE_ROCROLLER=OFFbuild can now select StreamK kernels it previously rejected, so kernel choice on that path changes; the numerical consequence of that selection change is measured later in this stack rather than here.Evaluating the predicates narrows that path further, and the measured blast radius is kernel choice rather than correctness: across the 8 shapes swept, no problem lost a solution and none failed verification, while 5 changed kernel. The cost is performance, not accuracy, and it falls on shapes that no checked-in rocRoller kernel was captured for — previously served by a kernel used outside its declared window. The largest measured case is about 3.6x slower. That cliff is removed by the fallback in #14497, later in this stack, not by this pull request.