[TRITON] Add gluon support for MXFP4 quant kernel in gfx950 and gfx1250#4193
[TRITON] Add gluon support for MXFP4 quant kernel in gfx950 and gfx1250#4193NimitPtl wants to merge 4 commits into
Conversation
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
NimitPtl
force-pushed
the
quant_mxfp4
branch
6 times, most recently
from
70ea8b9 to
dfea523
Compare
cagrikymk
changed the title
cagrikymk
changed the title
There was a problem hiding this comment.
Please move this file (aiter/ops/triton/gluon/quant_mxfp4.py) to aiter/ops/triton/_gluon_kernels/{gfx950,gfx1250}/quant/quant.py
_gluon_kernels folder is to store the arch specific gluon implementations. Also, given equivalent triton implementation lives under quant/quant.py , we can do the same for gluon implementations to mirror the existing structure.
| order=[0, 1], | ||
| ) | ||
|
|
||
| for pid_n in range(start_n, min(start_n + NUM_ITER, N), num_stages=NUM_STAGES): |
There was a problem hiding this comment.
Can you double check if num_stages has any effect for gluon implementations?
As far as I know the pass that creates multiple stages for the loops is specific to the triton backend:
If this is no-op, it will be confusing later.
If you want to add multi-stage support, you need to manually do so, similar to how it is done for the 1250 implementation (double buffering to break the dependency between memory op and quant op).
| ) | ||
| load_idx += 1 | ||
|
|
||
| x_reg = gl.amd.cdna4.async_copy.load_shared_relaxed( |
There was a problem hiding this comment.
I think it is better to use regular lds load here instead of cdna specific relaxed load.
| bs_offs_n = pid_n * NUM_QUANT_BLOCKS + gl.arange(0, NUM_QUANT_BLOCKS) | ||
| bs_offs = bs_offs_m[:, None] * stride_bs_m + bs_offs_n[None, :] * stride_bs_n | ||
| if EVEN_M_N: | ||
| gl.store(bs_ptr + bs_offs, bs_e8m0) |
There was a problem hiding this comment.
This can be added later, but TDM store might be useful here for perf.
2 participants
Motivation
Using gluon instead of triton for performance improvement
Technical Details
Gluon kernel, with sizes derived from the TTGIR of Triton kernel.
Test Plan
Test suite as well as benchmark script for testing accuracy and speed of quant kernel.
Test Result
Marginal improvement for gfx950 and around 10% improvement with usage of LDS in gfx1250
Submission Checklist