hotfix: fix the bug in #386 (#387)

Hotfix to #386, the input and output argument order in `SubtractLeft` and `FlagHeads` is different.
flashinfer-ai · Jul 21, 2024 · dc3f184 · dc3f184
1 parent 0cd4994
commit dc3f184
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/include/flashinfer/sampling.cuh b/include/flashinfer/sampling.cuh
@@ -115,7 +115,7 @@ __device__ __forceinline__ void DeviceSamplingFromProb(
     bool greater_than_u_diff[VEC_SIZE];
 #ifdef FLASHINFER_CUB_SUBTRACTLEFT_DEFINED
     BlockAdjacentDifference<bool, BLOCK_THREADS>(temp_storage->block_prim.adj_diff)
-        .SubtractLeft<VEC_SIZE>(greater_than_u_diff, greater_than_u, BoolDiffOp());
+        .SubtractLeft<VEC_SIZE>(greater_than_u, greater_than_u_diff, BoolDiffOp());
 #else
     BlockAdjacentDifference<bool, BLOCK_THREADS>(temp_storage->block_prim.adj_diff)
         .FlagHeads<VEC_SIZE>(greater_than_u_diff, greater_than_u, BoolDiffOp());