From 6bcfc2d97af65b205b649af5dcdfd30e1adfe50f Mon Sep 17 00:00:00 2001 From: Mark Saroufim Date: Tue, 9 Jun 2026 20:19:57 -0700 Subject: [PATCH] Use GEMM in L2 persistence CI test --- test/l2_persistence.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/test/l2_persistence.py b/test/l2_persistence.py index 266ec55..f749fe3 100644 --- a/test/l2_persistence.py +++ b/test/l2_persistence.py @@ -11,7 +11,7 @@ persisting lines (see ``csrc/clear_l2.cu``). On Hopper/Blackwell the default persisting carveout is already nonzero, so a submission only needs the access-policy window. -The test runs the SAME memory-bound GEMV through the isolated benchmark twice: +The test runs the SAME memory-bound single-column GEMM through the isolated benchmark twice: * ``kernel`` -- honest (each iteration reads the weights from a cleared L2). * ``kernel_l2_persist`` -- marks the reused weights persisting so they would survive the clear and stay warm across iterations. @@ -95,8 +95,16 @@ def _install_persistence(): ctypes.byref(val)) +def _gemm_vec(w, x, out=None): + x_col = x[:, None] + if out is None: + return torch.mm(w, x_col).squeeze(1).contiguous() + torch.mm(w, x_col, out=out[:, None]) + return out + + def kernel(output, x): - torch.mv(_get_state()["w"], x, out=output) + _gemm_vec(_get_state()["w"], x, out=output) def kernel_l2_persist(output, x): @@ -104,7 +112,7 @@ def kernel_l2_persist(output, x): if not _installed: _install_persistence() _installed = True - torch.mv(_get_state()["w"], x, out=output) + _gemm_vec(_get_state()["w"], x, out=output) def generate_test_case(*, seed): @@ -113,7 +121,7 @@ def generate_test_case(*, seed): x = torch.rand(_M, device="cuda", dtype=torch.float32, generator=gen).contiguous() w = _get_state()["w"] y = torch.empty(w.shape[0], device="cuda", dtype=torch.float32).contiguous() - expected = torch.mv(w, x).contiguous() + expected = _gemm_vec(w, x) return (y, x), (expected, 1e-2, 1e-2)