[fix] fix fa2a ops testcase

[duanjunwen]

Fix a bug involving mismatched parameter passing.

bash:
NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 torchrun --nnodes 1 --nproc_per_node 2 ./tests/test_fa2a_ops.py --world-size 2

error log:
Traceback (most recent call last): File "/root/duanjunwen/DistCA/./tests/test_fa2a_ops.py", line 262, in <module> Traceback (most recent call last): File "/root/duanjunwen/DistCA/./tests/test_fa2a_ops.py", line 262, in <module> test(args) File "/root/duanjunwen/DistCA/./tests/test_fa2a_ops.py", line 197, in test test(args) File "/root/duanjunwen/DistCA/./tests/test_fa2a_ops.py", line 197, in test a2a_memcpy_non_cp( TypeError: a2a_memcpy_non_cp() got an unexpected keyword argument 'copy_seq_mask' a2a_memcpy_non_cp( TypeError: a2a_memcpy_non_cp() got an unexpected keyword argument 'copy_seq_mask'
log after fix
copying rank 0 q shard to src buffer done. copying rank 0 q shard to src buffer done. copying rank 0 k shard to src buffer done. copying rank 0 k shard to src buffer done. copying rank 0 v shard to src buffer done. copying rank 0 v shard to src buffer done. copying rank 1 q shard to src buffer done. copying rank 1 q shard to src buffer done. copying rank 1 k shard to src buffer done. copying rank 1 k shard to src buffer done. copying rank 1 v shard to src buffer done. copying rank 1 v shard to src buffer done. copying rank 0 q shard from dst buffer done. copying rank 0 q shard from dst buffer done. copying rank 0 k shard from dst buffer done. copying rank 0 k shard from dst buffer done. copying rank 1 q shard from dst buffer done. copying rank 1 q shard from dst buffer done. copying rank 1 k shard from dst buffer done. copying rank 1 k shard from dst buffer done. copying rank 0 grads to send buffer done. copying rank 0 grads to send buffer done. copying rank 1 grads to send buffer done. copying rank 1 grads to send buffer done. copying rank 0 k shard from grad buffer done. copying rank 0 k shard from grad buffer done. copying rank 1 k shard from grad buffer done. copying rank 1 k shard from grad buffer done.

[GindaChen]

Thanks for the fix! We sure forgot to test out many of the unit tests after a refactor :( let me know if you hit others. Maybe this is a good time to let an agent do a pass...