From 6c92f9f328744901593683facfe2e5ad90a45104 Mon Sep 17 00:00:00 2001
From: Khoa Pham <khoa.pham@radixark.ai>
Date: Thu, 25 Jun 2026 18:13:57 -0700
Subject: [PATCH] fix(bench): pass DCP_RANK/DCP_WORLD_SIZE to
 set_mla_kv_buffer_kernel (#29356)

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
---
 test/registered/jit/benchmark/bench_set_mla_kv_buffer.py | 2 ++
 1 file changed, 2 insertions(+)

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,
     )