[Bugfix] Fix the incorrect offset computation for SGMV shrink (#22)

punica-ai · Nov 27, 2023 · 51cd092 · 51cd092
1 parent 91f1592
commit 51cd092
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/csrc/sgmv_flashinfer/sgmv_flashinfer.cuh b/csrc/sgmv_flashinfer/sgmv_flashinfer.cuh
@@ -69,7 +69,7 @@ __global__ void sgmv_shrink(T* y, T* x, T** w, IdType* s, float* tmp,
             offset += 8;
           }
           row_idx += 8;
-          y_ptr += 8 * d_out;
+          y_ptr += 8 * d_out - 2 * num_blocks_n * cell_capacity<T>();
           offset += 8 * num_cells_n - 4 * num_blocks_n;
         }
       }
@@ -329,7 +329,7 @@ __global__ void sgmv_shrink(T* y, T* x, T** w, IdType* s, float* tmp,
             offset += 8;
           }
           row_idx += 8;
-          y_ptr += 8 * d_out;
+          y_ptr += 8 * d_out - 2 * num_blocks_n * cell_capacity<T>();
           offset += 8 * num_cells_n - 4 * num_blocks_n;
         }
       }

diff --git a/tests/test_sgmv.py b/tests/test_sgmv.py
@@ -53,7 +53,7 @@ def lora_ref_impl(
         pytest.param("expand", marks=pytest.mark.xfail(reason="TODO: sgmv expand")),
     ],
 )
-@pytest.mark.parametrize("batch_setup", ["1x7", "7x1", "3x3"])
+@pytest.mark.parametrize("batch_setup", ["1x7", "7x1", "3x3", "32x1", "1x32"])
 @torch.inference_mode()
 def test_sgmv_correctness(dtype_str, h, r, direction, batch_setup):
     torch.manual_seed(0xABCDABCD987)
@@ -88,7 +88,7 @@ def test_sgmv_correctness(dtype_str, h, r, direction, batch_setup):
 
 @pytest.mark.xfail(reason="TODO: sgmv expand")
 @pytest.mark.parametrize("dtype_str", ["float16", "bfloat16"])
-@pytest.mark.parametrize("batch_setup", ["1x7", "7x1", "3x3"])
+@pytest.mark.parametrize("batch_setup", ["1x7", "7x1", "3x3", "32x1", "1x32"])
 @torch.inference_mode()
 def test_lora_correctness(dtype_str, batch_setup):
     torch.manual_seed(0xABCDABCD987)