misc:Remove unused k_smem_offset_w update in MLA kernel (#878)

muoshuosha · web-flow · commit 1605eaabff63 · 2025-02-19T10:13:50.000-08:00
The variable `ckv_smem_offset_w` and `kpe_smem_offset_w` are never used
after update in current loop, and they will be redefined and recomputed
in the next loop.I think these are redundant code.
diff --git a/include/flashinfer/attention/mla_fa2.cuh b/include/flashinfer/attention/mla_fa2.cuh
@@ -210,7 +210,6 @@ __device__ __forceinline__ void load_kv(
             warp_idx_in_wg * 4 + lane_idx / 8, 8 * mma_d + lane_idx % 8);
         ckv_smem.load_128b_async<SharedMemFillMode::kFillZero>(ckv_smem_offset_w, ckv_ptr,
                                                                q < kv_bound);
-        ckv_smem_offset_w = ckv_smem.template advance_offset_by_column<8>(ckv_smem_offset_w, mma_d);
         ckv_ptr += 8 * upcast_size<DTypeKV>();
       }
 
@@ -220,7 +219,6 @@ __device__ __forceinline__ void load_kv(
             warp_idx_in_wg * 4 + lane_idx / 8, 8 * mma_d + lane_idx % 8);
         kpe_smem.load_128b_async<SharedMemFillMode::kFillZero>(kpe_smem_offset_w, kpe_ptr,
                                                                q < kv_bound);
-        kpe_smem_offset_w = kpe_smem.template advance_offset_by_column<8>(kpe_smem_offset_w, mma_d);
         kpe_ptr += 8 * upcast_size<DTypeKV>();
       }
     }
@@ -245,7 +243,6 @@ __device__ __forceinline__ void load_kv(
             8 * mma_d + lane_idx % 8);
         ckv_smem.load_128b_async<SharedMemFillMode::kFillZero>(ckv_smem_offset_w, ckv_ptr,
                                                                q < kv_bound);
-        ckv_smem_offset_w = ckv_smem.template advance_offset_by_column<8>(ckv_smem_offset_w, mma_d);
         ckv_ptr += 8 * upcast_size<DTypeKV>();
       }
 
@@ -256,7 +253,6 @@ __device__ __forceinline__ void load_kv(
             8 * mma_d + lane_idx % 8);
         kpe_smem.load_128b_async<SharedMemFillMode::kFillZero>(kpe_smem_offset_w, kpe_ptr,
                                                                q < kv_bound);
-        kpe_smem_offset_w = kpe_smem.template advance_offset_by_column<8>(kpe_smem_offset_w, mma_d);
         kpe_ptr += 8 * upcast_size<DTypeKV>();
       }
     }

Original file line number	Diff line number	Diff line change
`@@ -210,7 +210,6 @@ __device__ __forceinline__ void load_kv(`
`210`	`210`	`warp_idx_in_wg * 4 + lane_idx / 8, 8 * mma_d + lane_idx % 8);`
`211`	`211`	`ckv_smem.load_128b_async<SharedMemFillMode::kFillZero>(ckv_smem_offset_w, ckv_ptr,`
`212`	`212`	`q < kv_bound);`
`213`		`- ckv_smem_offset_w = ckv_smem.template advance_offset_by_column<8>(ckv_smem_offset_w, mma_d);`
`214`	`213`	`ckv_ptr += 8 * upcast_size<DTypeKV>();`
`215`	`214`	`}`
`216`	`215`
`@@ -220,7 +219,6 @@ __device__ __forceinline__ void load_kv(`
`220`	`219`	`warp_idx_in_wg * 4 + lane_idx / 8, 8 * mma_d + lane_idx % 8);`
`221`	`220`	`kpe_smem.load_128b_async<SharedMemFillMode::kFillZero>(kpe_smem_offset_w, kpe_ptr,`
`222`	`221`	`q < kv_bound);`
`223`		`- kpe_smem_offset_w = kpe_smem.template advance_offset_by_column<8>(kpe_smem_offset_w, mma_d);`
`224`	`222`	`kpe_ptr += 8 * upcast_size<DTypeKV>();`
`225`	`223`	`}`
`226`	`224`	`}`
`@@ -245,7 +243,6 @@ __device__ __forceinline__ void load_kv(`
`245`	`243`	`8 * mma_d + lane_idx % 8);`
`246`	`244`	`ckv_smem.load_128b_async<SharedMemFillMode::kFillZero>(ckv_smem_offset_w, ckv_ptr,`
`247`	`245`	`q < kv_bound);`
`248`		`- ckv_smem_offset_w = ckv_smem.template advance_offset_by_column<8>(ckv_smem_offset_w, mma_d);`
`249`	`246`	`ckv_ptr += 8 * upcast_size<DTypeKV>();`
`250`	`247`	`}`
`251`	`248`
`@@ -256,7 +253,6 @@ __device__ __forceinline__ void load_kv(`
`256`	`253`	`8 * mma_d + lane_idx % 8);`
`257`	`254`	`kpe_smem.load_128b_async<SharedMemFillMode::kFillZero>(kpe_smem_offset_w, kpe_ptr,`
`258`	`255`	`q < kv_bound);`
`259`		`- kpe_smem_offset_w = kpe_smem.template advance_offset_by_column<8>(kpe_smem_offset_w, mma_d);`
`260`	`256`	`kpe_ptr += 8 * upcast_size<DTypeKV>();`
`261`	`257`	`}`
`262`	`258`	`}`