Skip to content

perf(fp8): Remove redundant clone calls from FP8 Linear Kernel #841

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
R3hankhan123 wants to merge 1 commit into
base: main
Choose a base branch
from
+17 −7
Open

perf(fp8): Remove redundant clone calls from FP8 Linear Kernel #841

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 10 additions & 4 deletions spyre_inference/custom_ops/fp8_linear_kernel.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Comment thread
R3hankhan123 marked this conversation as resolved.


def _activation_scale(x: torch.Tensor, per_token: bool) -> torch.Tensor:
Expand Down Expand Up @@ -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,
Expand Down
10 changes: 7 additions & 3 deletions tests/test_fp8_linear.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.

@github-actions github-actions Bot Sep 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These two cases (5, 130) do add M-padding, but the test still only asserts dtype and shape and only ever passes 2-D x. The bug this PR fixes is 3-D input that needs M-padding: the pre-clone() reshape reads the padded storage and returns wrong numbers. On device I confirmed it — feeding the same 130 rows as (13, 10, 128) gives max|out3d - out2d| = 2.999 on main and 0.0 on this branch. Neither of those facts is checked here: no 3-D shape is exercised (so the reshape branch never runs in CI), and no value comparison is done (so even a 2-D numeric regression would pass). The docstring's claim that these cases are "verifying the trim-before-reshape path" is therefore not accurate. Please add a 3-D case and compare apply_weights output against a reference (e.g. the 2-D result reshaped, or a dequantized fp16 matmul), as tdoublep asked in the review.

@joerunde joerunde Sep 11, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@R3hankhan123 looks like this might need one more pass then we should be good to go

"""
if not spyre_available():
pytest.skip("Spyre device not available")
if SpyreFp8LinearKernel is None:
Expand All @@ -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():
Expand Down
Loading