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>
2026-06-25 18:13:57 -07:00
parent 781537b61d
commit 6c92f9f328
1 changed files with 2 additions and 0 deletions
@@ -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,
    )