Repository navigation
Conversation
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
There was a problem hiding this comment.
Pull request overview
Adds a missing Triton test module so the bench_gemm_afp4wfp4_pre_quant_atomic.py benchmark can import a shared input-generator helper without failing with ModuleNotFoundError.
Changes:
- Added
op_tests/triton_tests/gemm/basic/test_gemm_afp4wfp4_pre_quant_atomic.py. - Implemented
generate_gemm_afp4wfp4_pre_quant_inputs(...)to create(x, w, None, w_scales, y)inputs forgemm_afp4wfp4_pre_quant. - Added a basic pytest invocation that runs the kernel and checks output shape/dtype.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| if layout[1] == "N": | ||
| w = torch.randint(0, 256, (N, K // 2), dtype=torch.uint8, device="cuda") | ||
| else: | ||
| w = torch.randint(0, 256, (K // 2, N), dtype=torch.uint8, device="cuda").T | ||
|
|
||
| w_scales = torch.randint( | ||
| 124, 128, (N, K // SCALE_GROUP_SIZE), dtype=torch.uint8, device="cuda" | ||
| ) |
There was a problem hiding this comment.
K is used with floor division for both the packed-weight dimension (K // 2) and the scale dimension (K // SCALE_GROUP_SIZE). If a caller passes a K that is not divisible by 2 and 32 respectively, this will silently generate tensors with incompatible shapes (and likely lead to incorrect results or hard-to-debug kernel errors). Add explicit validation (e.g., assert/ValueError) that K % 2 == 0 and K % SCALE_GROUP_SIZE == 0 (and ideally also that SCALE_GROUP_SIZE matches the kernel’s expectation).
| y = None | ||
| if output: | ||
| y = torch.zeros((M, N), dtype=torch.bfloat16, device="cuda") | ||
|
|
||
| return x, w, None, w_scales, y |
There was a problem hiding this comment.
The helper always preallocates y as BF16 when output=True, but gemm_afp4wfp4_pre_quant(..., dtype=...) only uses the dtype argument when it allocates y internally; if a preallocated y is provided, the output dtype is effectively y.dtype (see aiter/ops/triton/gemm/basic/gemm_a16wfp4.py where the kernel stores to c_ptr.type.element_ty). This makes it easy for callers (e.g., the benchmark) to pass dtype=torch.float32 but still get BF16 output. Consider adding a y_dtype/out_dtype parameter (or taking dtype and using it for y allocation) so the helper can produce a y that matches the requested output dtype.
| result = gemm_afp4wfp4_pre_quant(x, w, w_scales, c_dtype, y) | ||
| assert result.shape == (M, N) | ||
| assert result.dtype == c_dtype |
There was a problem hiding this comment.
This test only asserts shape and dtype, so it won’t catch regressions where the kernel returns numerically incorrect results (especially important given FP4 + scaling). Other GEMM tests in this directory validate against a torch reference (e.g., test_gemm_a16wfp4.py). Consider adding at least one small-shape correctness check (or reusing existing FP4 dequant helpers) to make this a functional test rather than just a smoke test.
|
This is a deprecated kernel wrapper - its superseded by gemm_a16wfp4 and will be removed. The solution is to delete the bench as well. cc @vgokhale |
The benchmark cannot run: it imports generate_gemm_afp4wfp4_pre_quant_inputs from op_tests.triton_tests.gemm.basic.test_gemm_afp4wfp4_pre_quant_atomic, a module that no longer exists. The DS FP4 fusion work (#1371, #1434) renamed that test to test_gemm_a16wfp4.py, so every run fails with ModuleNotFoundError. #2601 tried to restore the helper. It was closed because gemm_afp4wfp4_pre_quant is a deprecated wrapper, superseded by gemm_a16wfp4 and slated for removal, and the bench should be deleted. This does that, and drops the bench's entry from bench_schema.yaml.
Motivation
bench_gemm_afp4wfp4_pre_quant_atomic.pyfails withModuleNotFoundErrorbecause it importedgenerate_gemm_afp4wfp4_pre_quant_inputsfromop_tests.triton_tests.gemm.basic.test_gemm_afp4wfp4_pre_quant_atomic, a test file that did notexist.
Technical Details
Added
op_tests/triton_tests/gemm/basic/test_gemm_afp4wfp4_pre_quant_atomic.pywith agenerate_gemm_afp4wfp4_pre_quant_inputshelper. Sincegemm_afp4wfp4_pre_quantis a wrapper aroundgemm_a16wfp4, activations are BF16 (quantizedon-the-fly by the kernel) rather than packed FP4. The generator produces (x, w, None, w_scales, y) to match the benchmark's unpacking convention.
Test Plan
Run
bench_gemm_afp4wfp4_pre_quant_atomic.pywith a representative shape and verify it no longer raisesModuleNotFoundError.Test Result
The benchmark script launches successfully after the fixes.