feat: support mlu mla pd pull with mooncake transfer. by phantomlei3 · Pull Request #1167 · jd-opensource/xllm

phantomlei3 · 2026-04-02T11:44:11Z

No description provided.

gemini-code-assist

Code Review

This pull request enables Disaggregated Prefill/Decode (PD) support for the MLU backend by integrating and refactoring the Mooncake transfer engine. Key changes include the introduction of a descriptor-based memory registration system (using BufferDesc and MemInfo) and a new GetRegisteredMemory RPC to synchronize memory layouts between nodes, which is essential for supporting complex architectures like MLA. Additionally, the PR enhances validation logic in LLMEngine::pull_kv_blocks, automates device_ip resolution, and adds comprehensive unit tests for the transfer engine. Review feedback identifies several opportunities to align with the project's style guide, specifically regarding the use of emplace_back for vector operations and the annotation of constant arguments in function calls.

gemini-code-assist · 2026-04-02T11:50:57Z

xllm/core/framework/kv_cache/mooncake_transfer_engine.cpp

  buffers.reserve(num);
  for (size_t i = 0; i < num; i++) {
-    buffers.push_back(BufferEntry{(void*)addrs[i], lens[i]});
+    buffers.push_back(BufferEntry{addrs[i], lens[i]});


According to the style guide (rule 291), emplace_back should be preferred over push_back to construct elements in-place and avoid unnecessary copies.

Suggested change

buffers.push_back(BufferEntry{addrs[i], lens[i]});

buffers.emplace_back(addrs[i], lens[i]);

References

Rule 291 states to prefer emplace_back over push_back for in-place construction and to avoid copies. ^(link)

gemini-code-assist · 2026-04-02T11:50:57Z

xllm/core/framework/kv_cache/mooncake_transfer_engine.cpp

+    entries.push_back(BufferEntry{
+        reinterpret_cast<void*>(buffer_desc.addr),
+        static_cast<size_t>(buffer_desc.len),
+    });


Per style guide rule 291, please use emplace_back instead of push_back to construct elements in-place.

entries.emplace_back(reinterpret_cast<void*>(buffer_desc.addr), static_cast<size_t>(buffer_desc.len));

References

Rule 291 states to prefer emplace_back over push_back for in-place construction and to avoid copies. ^(link)

gemini-code-assist · 2026-04-02T11:50:57Z

xllm/core/framework/kv_cache/mooncake_transfer_engine.cpp

+        TransferRequest entry;
+        entry.opcode = opcode;
+        entry.length = len;
+        entry.source = static_cast<void*>(local_base + local_bias);
+        entry.target_id = remote_handle;
+        entry.target_offset =
+            reinterpret_cast<uint64_t>(remote_base + remote_bias);
+        entry.advise_retry_cnt = 0;
+        entries->push_back(entry);


According to the style guide (rule 291), emplace_back should be preferred over push_back to construct elements in-place and avoid unnecessary copies. You can construct the TransferRequest directly in the vector.

Suggested change

TransferRequest entry;

entry.opcode = opcode;

entry.length = len;

entry.source = static_cast<void*>(local_base + local_bias);

entry.target_id = remote_handle;

entry.target_offset =

reinterpret_cast<uint64_t>(remote_base + remote_bias);

entry.advise_retry_cnt = 0;

entries->push_back(entry);

TransferRequest& entry = entries->emplace_back();

entry.opcode = opcode;

entry.length = len;

entry.source = static_cast<void*>(local_base + local_bias);

entry.target_id = remote_handle;

entry.target_offset =

reinterpret_cast<uint64_t>(remote_base + remote_bias);

entry.advise_retry_cnt = 0;

References

Rule 291 states to prefer emplace_back over push_back for in-place construction and to avoid copies. ^(link)

gemini-code-assist · 2026-04-02T11:50:57Z

xllm/core/framework/kv_cache/mooncake_transfer_engine.cpp

    TransferRequest entry;
    entry.opcode = opcode;
    entry.length = transfer_size;
-    entry.source = (void*)(local_base + src_offsets[i]);
+    entry.source = (void*)(local_base + local_offset);
    entry.target_id = remote_handle;
-    entry.target_offset = (uint64_t)(remote_base + dst_offsets[i]);
+    entry.target_offset = (uint64_t)(remote_base + remote_offset);
    entry.advise_retry_cnt = 0;
    entries.push_back(entry);


According to the style guide (rule 291), emplace_back should be preferred over push_back to construct elements in-place and avoid unnecessary copies. You can construct the TransferRequest directly in the vector.

Suggested change

TransferRequest entry;

entry.opcode = opcode;

entry.length = transfer_size;

entry.source = (void*)(local_base + src_offsets[i]);

entry.source = (void*)(local_base + local_offset);

entry.target_id = remote_handle;

entry.target_offset = (uint64_t)(remote_base + dst_offsets[i]);

entry.target_offset = (uint64_t)(remote_base + remote_offset);

entry.advise_retry_cnt = 0;

entries.push_back(entry);

TransferRequest& entry = entries.emplace_back();

entry.opcode = opcode;

entry.length = transfer_size;

entry.source = (void*)(local_base + local_offset);

entry.target_id = remote_handle;

entry.target_offset = (uint64_t)(remote_base + remote_offset);

entry.advise_retry_cnt = 0;

References

Rule 291 states to prefer emplace_back over push_back for in-place construction and to avoid copies. ^(link)

gemini-code-assist · 2026-04-02T11:50:58Z

xllm/core/framework/kv_cache/mooncake_transfer_engine_test.cpp

+  local_info.layers[0].key =
+      mk_desc(0, proto::BUFFER_KIND_KEY, 0x1000, 0x400, 0x40);
+  local_info.layers[0].index =
+      mk_desc(0, proto::BUFFER_KIND_INDEX, 0x3000, 0x80, 0x08);
+  remote_info.layers[0].key =
+      mk_desc(0, proto::BUFFER_KIND_KEY, 0x5000, 0x400, 0x40);
+  remote_info.layers[0].index =
+      mk_desc(0, proto::BUFFER_KIND_INDEX, 0x7000, 0x80, 0x08);


According to the style guide (rule 260), constant arguments in function calls should be annotated with comments indicating the parameter name. This improves readability, especially for functions with multiple arguments of similar types. Please add annotations to the calls to mk_desc.

Suggested change

local_info.layers[0].key =

mk_desc(0, proto::BUFFER_KIND_KEY, 0x1000, 0x400, 0x40);

local_info.layers[0].index =

mk_desc(0, proto::BUFFER_KIND_INDEX, 0x3000, 0x80, 0x08);

remote_info.layers[0].key =

mk_desc(0, proto::BUFFER_KIND_KEY, 0x5000, 0x400, 0x40);

remote_info.layers[0].index =

mk_desc(0, proto::BUFFER_KIND_INDEX, 0x7000, 0x80, 0x08);

local_info.layers[0].key =

mk_desc(/*layer_id=*/0, /*kind=*/proto::BUFFER_KIND_KEY, /*addr=*/0x1000, /*len=*/0x400, /*bytes_per_block=*/0x40);

local_info.layers[0].index =

mk_desc(/*layer_id=*/0, /*kind=*/proto::BUFFER_KIND_INDEX, /*addr=*/0x3000, /*len=*/0x80, /*bytes_per_block=*/0x08);

remote_info.layers[0].key =

mk_desc(/*layer_id=*/0, /*kind=*/proto::BUFFER_KIND_KEY, /*addr=*/0x5000, /*len=*/0x400, /*bytes_per_block=*/0x40);

remote_info.layers[0].index =

mk_desc(/*layer_id=*/0, /*kind=*/proto::BUFFER_KIND_INDEX, /*addr=*/0x7000, /*len=*/0x80, /*bytes_per_block=*/0x08);

References

Rule 260 requires annotating constant arguments in function calls with comments indicating the parameter name to improve readability. ^(link)

gemini-code-assist · 2026-04-02T11:50:58Z

xllm/core/framework/kv_cache/mooncake_transfer_engine_test.cpp

+  local_info.layers[0].key =
+      mk_desc(0, proto::BUFFER_KIND_KEY, 0x1000, 0x400, 0x40);
+  local_info.layers[1].key =
+      mk_desc(1, proto::BUFFER_KIND_KEY, 0x2000, 0x400, 0x40);
+  local_info.layers[1].index =
+      mk_desc(1, proto::BUFFER_KIND_INDEX, 0x3000, 0x80, 0x08);
+  remote_info.layers[0].key =
+      mk_desc(0, proto::BUFFER_KIND_KEY, 0x5000, 0x400, 0x40);
+  remote_info.layers[1].key =
+      mk_desc(1, proto::BUFFER_KIND_KEY, 0x6000, 0x400, 0x40);
+  remote_info.layers[1].index =
+      mk_desc(1, proto::BUFFER_KIND_INDEX, 0x7000, 0x80, 0x08);


The calls to mk_desc are missing annotations for constant arguments, which is required by the style guide (rule 260) to improve code readability.

Suggested change

local_info.layers[0].key =

mk_desc(0, proto::BUFFER_KIND_KEY, 0x1000, 0x400, 0x40);

local_info.layers[1].key =

mk_desc(1, proto::BUFFER_KIND_KEY, 0x2000, 0x400, 0x40);

local_info.layers[1].index =

mk_desc(1, proto::BUFFER_KIND_INDEX, 0x3000, 0x80, 0x08);

remote_info.layers[0].key =

mk_desc(0, proto::BUFFER_KIND_KEY, 0x5000, 0x400, 0x40);

remote_info.layers[1].key =

mk_desc(1, proto::BUFFER_KIND_KEY, 0x6000, 0x400, 0x40);

remote_info.layers[1].index =

mk_desc(1, proto::BUFFER_KIND_INDEX, 0x7000, 0x80, 0x08);

local_info.layers[0].key =

mk_desc(/*layer_id=*/0, /*kind=*/proto::BUFFER_KIND_KEY, /*addr=*/0x1000, /*len=*/0x400, /*bytes_per_block=*/0x40);

local_info.layers[1].key =

mk_desc(/*layer_id=*/1, /*kind=*/proto::BUFFER_KIND_KEY, /*addr=*/0x2000, /*len=*/0x400, /*bytes_per_block=*/0x40);

local_info.layers[1].index =

mk_desc(/*layer_id=*/1, /*kind=*/proto::BUFFER_KIND_INDEX, /*addr=*/0x3000, /*len=*/0x80, /*bytes_per_block=*/0x08);

remote_info.layers[0].key =

mk_desc(/*layer_id=*/0, /*kind=*/proto::BUFFER_KIND_KEY, /*addr=*/0x5000, /*len=*/0x400, /*bytes_per_block=*/0x40);

remote_info.layers[1].key =

mk_desc(/*layer_id=*/1, /*kind=*/proto::BUFFER_KIND_KEY, /*addr=*/0x6000, /*len=*/0x400, /*bytes_per_block=*/0x40);

remote_info.layers[1].index =

mk_desc(/*layer_id=*/1, /*kind=*/proto::BUFFER_KIND_INDEX, /*addr=*/0x7000, /*len=*/0x80, /*bytes_per_block=*/0x08);

References

Rule 260 requires annotating constant arguments in function calls with comments indicating the parameter name to improve readability. ^(link)

feat: support mlu mla pd pull with mooncake transfer.

c8d119d

phantomlei3 requested review from DongheJin, JimHsiung, RobbieLeung, XuZhang99, liutongxuan, walsonyang and yq33victor as code owners April 2, 2026 11:44

phantomlei3 requested a review from sunbaosong April 2, 2026 11:44

gemini-code-assist bot reviewed Apr 2, 2026

View reviewed changes

JimHsiung approved these changes Apr 3, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: support mlu mla pd pull with mooncake transfer.#1167

feat: support mlu mla pd pull with mooncake transfer.#1167
phantomlei3 wants to merge 1 commit intojd-opensource:mainfrom
phantomlei3:feat/pd-support-mlu

phantomlei3 commented Apr 2, 2026

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot Apr 2, 2026

Uh oh!

gemini-code-assist bot Apr 2, 2026

Uh oh!

gemini-code-assist bot Apr 2, 2026

Uh oh!

gemini-code-assist bot Apr 2, 2026

Uh oh!

gemini-code-assist bot Apr 2, 2026

Uh oh!

gemini-code-assist bot Apr 2, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

	buffers.push_back(BufferEntry{addrs[i], lens[i]});
	buffers.emplace_back(addrs[i], lens[i]);

Conversation

phantomlei3 commented Apr 2, 2026

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Apr 2, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Apr 2, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Apr 2, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Apr 2, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Apr 2, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Apr 2, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants