diff --git a/test/registered/jit/benchmark/bench_set_mla_kv_buffer.py b/test/registered/jit/benchmark/bench_set_mla_kv_buffer.py
index 58a0d9d9a..71b4d20f7 100644
--- a/test/registered/jit/benchmark/bench_set_mla_kv_buffer.py
+++ b/test/registered/jit/benchmark/bench_set_mla_kv_buffer.py
@@ -48,6 +48,8 @@ def _triton_baseline(kv_buffer, loc, cache_k_nope, cache_k_rope):
         nope_dim,
         rope_dim,
         BLOCK=BLOCK,
+        DCP_RANK=0,
+        DCP_WORLD_SIZE=1,
         **pdl_kwargs,
     )