[Bugfix][PD] Accept slot_layer_ids in NixlKVManager.set_kv_buffer_tensors

[he-yufeng]

Motivation

Fixes #42684.

With SGLANG_DISAGG_STAGING_BUFFER=1 on non-MLA models, both prefill.py and decode.py call

kv_manager.set_kv_buffer_tensors(k_buffers, v_buffers, page_size, slot_layer_ids=...)

unconditionally. MooncakeKVManager accepts the slot_layer_ids kwarg, but NixlKVManager never got it, so picking the NIXL backend in staging mode kills the kv manager with

TypeError: NixlKVManager.set_kv_buffer_tensors() got an unexpected keyword argument 'slot_layer_ids'

during init and the server never comes up.

Modifications

NixlKVManager.set_kv_buffer_tensors now takes slot_layer_ids and stashes it in kv_buffer_tensors, matching the Mooncake signature. This is parity only: the NIXL staged gather walks the registered buffers in list order and nothing reads the mapping back. The same scoped fix is currently carried inside the larger #40909 and #40861; this PR lands it on its own so staging mode works on main without waiting for those.

Added two unit tests to TestNixlStaging covering the kwarg and its default.

Accuracy Tests

No model output changes; disaggregation metadata path only.

Speed Tests and Profiling

No hot path impact.

Checklist

• Format your code according to the Format code with pre-commit. (ruff check --select=F401,F821,F823,UP037 and ruff format --check both pass on the touched files)
• Add unit tests according to the Run and add unit tests. The two new cases live in TestNixlStaging (base-a-test-cpu); I do not have NIXL hardware here, so the staging transfer paths themselves are covered by CI, as before.
• Update documentation according to Write documentations. (N/A, no user-facing config change)
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed. (N/A, see above)
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ❌ Run #37628033810
Latest PR Test (Extra): ❌ Run #37628033368
Latest PR Test (AMD ROCm 10): ❌ Run #37628033989

[he-yufeng]

One honest note on verification: I tried to run test/registered/unit/disaggregation/test_nixl_backend_basic.py locally before pushing, but the run kept hanging in the import chain in my environment and never produced a result, so I would rather lean on CI here than claim a local pass I do not have. Could someone with write access add the run-ci label? The pr-gate job fails fast on the missing label, so nothing has actually executed the test yet.

[he-yufeng]

Closing as superseded by #43723 ("[PD][NIXL] Fix staging buffer startup by accepting slot_layer_ids"), which landed the same fix on main: NixlKVManager.set_kv_buffer_tensors now accepts slot_layer_ids and stashes it the same way, and test_staging_draft_kv_slots.py::test_every_backend_setter_accepts_slot_layer_ids covers the acceptance contract across every backend, which is broader than the two NIXL-only tests here. Nothing remains to port.