From 8b7bc81a5d7b4a4a92b0f259398d58afd3749d6b Mon Sep 17 00:00:00 2001 From: Rehan Khan Date: Thu, 10 Sep 2026 17:52:35 +0530 Subject: [PATCH] perf(fp8): Remove redundant clone calls Signed-off-by: Rehan Khan --- spyre_inference/custom_ops/fp8_linear_kernel.py | 14 ++++++++++---- tests/test_fp8_linear.py | 10 +++++++--- 2 files changed, 17 insertions(+), 7 deletions(-) diff --git a/spyre_inference/custom_ops/fp8_linear_kernel.py b/spyre_inference/custom_ops/fp8_linear_kernel.py index 8ec129640..a899918b0 100644 --- a/spyre_inference/custom_ops/fp8_linear_kernel.py +++ b/spyre_inference/custom_ops/fp8_linear_kernel.py @@ -76,8 +76,8 @@ def _n_tiles(n: int) -> list[int]: def _join(parts: list[torch.Tensor], dim: int) -> torch.Tensor: - """Cat tiles into a new buffer so RMSNorm/SiLU/attention see offset 0.""" - return (parts[0] if len(parts) == 1 else torch.cat(parts, dim=dim)).clone() + """Concatenate tiles; both ``_fp8_mm`` and ``cat`` already return fresh buffers.""" + return parts[0] if len(parts) == 1 else torch.cat(parts, dim=dim) def _activation_scale(x: torch.Tensor, per_token: bool) -> torch.Tensor: @@ -261,10 +261,16 @@ def apply_weights( col_outs.append(_fp8_mm(xi, wj, sj, bj, self._per_token_act)) col += ns row_outs.append(_join(col_outs, dim=-1)) - out = _join(row_outs, dim=0)[:orig_m] + out = _join(row_outs, dim=0) + if out.shape[0] > orig_m: + # The slice is already contiguous at offset 0; clone() compacts the + # storage so the subsequent reshape (3-D inputs) sees the correct + # element count without it the padding rows corrupt the trailing + # dimensions. + out = out[:orig_m].clone() if x.dim() > 2: out = out.reshape(*orig_shape[:-1], out.shape[-1]) - return out.clone() + return out def apply_scaled_mm( self, diff --git a/tests/test_fp8_linear.py b/tests/test_fp8_linear.py index 23ecab193..59552cf77 100644 --- a/tests/test_fp8_linear.py +++ b/tests/test_fp8_linear.py @@ -204,9 +204,13 @@ def _run_spyre_apply(self, kernel, layer, x): assert actual.device.type == "spyre", actual.device return actual - @pytest.mark.parametrize("num_tokens", [1, 4, 128]) + @pytest.mark.parametrize("num_tokens", [1, 4, 5, 128, 130]) def test_scaled_mm_apply(self, num_tokens): - """apply_weights runs aten._scaled_mm on Spyre.""" + """apply_weights runs aten._scaled_mm on Spyre. + + num_tokens=5 and 130 exercise M-padding (not in _SMALL_M, not aligned + to _M_ALIGN=128), verifying the trim-before-reshape path. + """ if not spyre_available(): pytest.skip("Spyre device not available") if SpyreFp8LinearKernel is None: @@ -228,7 +232,7 @@ def test_scaled_mm_apply(self, num_tokens): assert actual.dtype == torch.float16 assert actual.shape == (num_tokens, out_features) - @pytest.mark.parametrize("num_tokens", [1, 4, 128]) + @pytest.mark.parametrize("num_tokens", [1, 4, 5, 128, 130]) def test_scaled_mm_apply_per_channel(self, num_tokens): """apply_weights with Granite per-channel weight scales + per-token acts.""" if not spyre_available():