[proof-of-concept] add MXFP8 pre-swizzling for gfx1250#568
[proof-of-concept] add MXFP8 pre-swizzling for gfx1250#568matthiasdiener wants to merge 18 commits intodevfrom
Conversation
matthiasdiener
force-pushed
the
mdiener/mxfp8-swizzle
branch
from
ddf19da to
313a6b7
Compare
| asm volatile("ds_swizzle_b32 %0, %1 offset:0x041F\n\t" | ||
| "s_waitcnt lgkmcnt(0)" : "=v"(r) : "v"(v)); | ||
| return r; | ||
| return __shfl_xor(v, 1); |
There was a problem hiding this comment.
Do we still need these helper functions now that we're just doing a __shfl_xor?
There was a problem hiding this comment.
This change is only inadvertently part of this PR, it is already part of #571. Will revert here.
| // Col-wise: input is [K_scale, M] row-major (M contiguous), representing | ||
| // the column-wise scale matrix logically shaped [M, K_scale]. | ||
| // Logical (m, k) maps to physical address k * original_M + m. | ||
| __global__ void __launch_bounds__(256) |
There was a problem hiding this comment.
This function is almost identical to the rowwise scaling func, can we merge them and template? Ideally we could have 1 thread write both the colwise and rowwise scale if we are doing both.
| const int k = idx % K_scale; | ||
|
|
||
| uint8_t val = 127; | ||
| if (m < original_M && k < original_K) { |
There was a problem hiding this comment.
Could we move this check to the hostside, or remove it completely?
| : public ::testing::TestWithParam< | ||
| std::tuple<std::pair<int, int>, bool>> {}; | ||
|
|
||
| TEST_P(MxSwizzleTestSuite, TestMxSwizzle) { |
There was a problem hiding this comment.
I think full GEMM tests should live in test_cublaslt.cu. Also, wondering if this is not already covered by the MXFP8 gemm tests present there, if we are always swizzling. Probably should limit tests here to just testing the swizzled scales, if needed at all.
| #include <cstdint> | ||
|
|
||
| #include "../common.h" | ||
| #include "../util/cuda_runtime.h" |
There was a problem hiding this comment.
Why is this include needed?
| " (got shape=", shape, ")"); | ||
| #ifdef USE_ROCM | ||
| // gfx1250 MX pre-swizzle (Tensile 3D) layout requires M padded to multiple of 4. | ||
| // Other ROCm architectures use 128x4 tiles but currently skip padding |
There was a problem hiding this comment.
I'm not sure this is true regarding us using 128x4 tiles. 128x4 scaling is an upstream requirement. We also have padding expectations in pytorch, jax, and all 3 test dirs have padding that will probably need fixing.
2 participants
Description
Please include a brief summary of the changes, relevant motivation and context.
Fixes https://github.com/ROCm/frameworks-internal/issues/16428
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: