[ROCm][CI] Pass weight_shape in MXFP8 block32 linear tests (#58698)
Signed-off-by: Djordje Ramic <djoramic@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
djramic committed
Sep 25, 2026 at 12:52 UTC
25b0add7b8a1c944d5c4e364f2de6aa82497a2ad
1 file changed
+5
-4
tests/kernels/quantization/test_rocm_mxfp8_linear.py
+5
-4
@@ -166,7 +166,7 @@ def test_rocm_mxfp8_linear_block32_matches_dequant_reference(name, n, k, num_tok
166
device = "cuda"
167
layer = _make_block32_layer(n, k, device)
168
row_scale = layer.weight_scale.data.clone()
169
- kernel = RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig(bmm_batch_size=None))
169
+ kernel = RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig(weight_shape=(n, k)))
170
kernel.process_weights_after_loading(layer)
171
assert layer.weight.dtype == torch.float8_e4m3fn
172
assert layer.weight_scale.shape == (n // 32, k // 32)
@@ -194,11 +194,12 @@ def test_rocm_mxfp8_linear_keeps_per_row_scales():
194
RocmDotScaledMxfp8LinearKernel,
195
)
196
197
- layer = _make_block32_layer(64, 256, "cuda")
197
+ n, k = 64, 256
198
+ layer = _make_block32_layer(n, k, "cuda")
199
layer.weight_scale.data[1, 0] += 1
199
- kernel = RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig(bmm_batch_size=None))
200
+ kernel = RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig(weight_shape=(n, k)))
201
kernel.process_weights_after_loading(layer)
201
- assert layer.weight_scale.shape == (64, 256 // 32)
202
+ assert layer.weight_scale.shape == (n, k // 32)
203
204
205
def _split_k_case(packed: bool):