From f2170fc1b3597ec81c4835845af1c989b3ed367f Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 24 May 2026 06:41:25 +0000 Subject: [PATCH] =?UTF-8?q?D1:=20Fix=20kvb=E2=86=92kvh=20typo=20in=20PV=20?= =?UTF-8?q?GEMM?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- dsv4/kernels/attention/fmha.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index 8a0bee2e..5d8ca0b2 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -263,11 +263,11 @@ class FmhaKernel: pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0) if not self.use_smem_p: for kb in cutlass.range(cute.size(tOrP0, mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kvb.index)], tOtO0) + cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,kvh.index)], tOtO0) pv_mma.set(tcgen05.Field.ACCUMULATE, True) else: for kb in cutlass.range(cute.size(tCrP, mode=[2]), unroll_full=True): - cute.gemm(pv_mma, tOtO0, tCrP[(None,None,kb,0)], tCrV[(None,None,kvb.index)], tOtO0) + cute.gemm(pv_mma, tOtO0, tCrP[(None,None,kb,0)], tCrV[(None,None,kb,kvh.index)], tOtO0) pv_mma.set(tcgen05.Field.ACCUMULATE, True) cute.arch.fence_view_async_tmem_store() final_o_bar.arrive()