【WIP】feat: onerec xattn npu multistream. by DragonFive · Pull Request #1453 · jd-opensource/xllm

DragonFive · 2026-05-14T07:59:21Z

No description provided.

DragonFive · 2026-05-14T08:01:07Z

 OneRecBatchInputBuilder::HighPerformanceCache&
 OneRecBatchInputBuilder::get_perf_cache() {
-  static HighPerformanceCache cache;
+  thread_local HighPerformanceCache cache;


here no need to change, because we won't use multi stream in prefill only, and this cache only works in prefill only mode.

DragonFive · 2026-05-14T08:02:41Z

+        {static_cast<int64_t>(workspace_size)},
+        torch::TensorOptions().dtype(torch::kUInt8).device(logprobs.device()));
+    workspace_addr = workspace_tensor.data_ptr();
+    c10_npu::NPUCachingAllocator::recordStream(


why you delete below aclrtSynchronizeStream and only add recordStream?

DragonFive · 2026-05-14T08:03:15Z

+        torch::TensorOptions().dtype(torch::kUInt8).device(logprobs.device()));
+    workspace_addr = workspace_tensor.data_ptr();
+    c10_npu::NPUCachingAllocator::recordStream(
+        workspace_tensor.storage().data_ptr(), npu_stream);


what's this means for workspace_tensor and npu_stream in recordstream?

gemini-code-assist

Code Review

This pull request introduces support for multi-concurrency in the OneRec xattention pipeline, enabling concurrent batch preparation and execution. Key changes include transitioning the performance cache to thread-local storage, implementing serialization for NPU-specific preparation and forward passes to ensure thread safety, and updating the scheduler to manage in-flight steps and resource exhaustion more effectively. Additionally, NPU kernel workspace management was improved by using the NPUCachingAllocator instead of manual allocation. Feedback includes several requests to annotate constant arguments with parameter names to comply with the repository style guide, a critical recommendation to capture asynchronous lambda parameters by value to avoid dangling references, and a suggestion to add a null check for the ATB context before use.

gemini-code-assist · 2026-05-14T08:03:20Z

+  FLAGS_max_decode_rounds = 2;
+  FLAGS_prefill_scheduling_memory_usage_threshold = 1.0;
+
+  std::unique_ptr<FakeEngine> engine = std::make_unique<FakeEngine>(64, 128);


Constant arguments in function or constructor calls should be annotated with a comment indicating the parameter name to improve readability, as per the repository style guide.

Suggested change

std::unique_ptr<FakeEngine> engine = std::make_unique<FakeEngine>(64, 128);

std::unique_ptr<FakeEngine> engine = std::make_unique<FakeEngine>(/*num_blocks=*/64, /*block_size=*/128);

References

Annotate constant arguments with a comment indicating the parameter name when calling functions or constructors. ^(link)

gemini-code-assist · 2026-05-14T08:03:20Z

+      /*rec_worker_max_concurrency=*/2);
+  FixedStepsScheduler scheduler(engine.get(), opt);
+  std::vector<std::shared_ptr<Request>> requests =
+      GenRequests({1, 1}, {1, 1}, RecType::kOneRec);


Constant arguments in function calls should be annotated with a comment indicating the parameter name to improve readability.

Suggested change

GenRequests({1, 1}, {1, 1}, RecType::kOneRec);

GenRequests(/*prompt_lens=*/{1, 1}, /*max_tokens=*/{1, 1}, RecType::kOneRec);

References

Annotate constant arguments with a comment indicating the parameter name when calling functions or constructors. ^(link)

gemini-code-assist · 2026-05-14T08:03:20Z

+  FLAGS_max_decode_rounds = 2;
+  FLAGS_prefill_scheduling_memory_usage_threshold = 1.0;
+
+  std::unique_ptr<FakeEngine> engine = std::make_unique<FakeEngine>(3, 128);


Constant arguments in function or constructor calls should be annotated with a comment indicating the parameter name.

Suggested change

std::unique_ptr<FakeEngine> engine = std::make_unique<FakeEngine>(3, 128);

std::unique_ptr<FakeEngine> engine = std::make_unique<FakeEngine>(/*num_blocks=*/3, /*block_size=*/128);

References

Annotate constant arguments with a comment indicating the parameter name when calling functions or constructors. ^(link)

gemini-code-assist · 2026-05-14T08:03:20Z

+  FLAGS_prefill_scheduling_memory_usage_threshold = 1.0;
+
+  std::unique_ptr<FakeEngine> engine = std::make_unique<FakeEngine>(3, 128);
+  engine->notify_after_step_calls(2);


Constant arguments in function calls should be annotated with a comment indicating the parameter name.

Suggested change

engine->notify_after_step_calls(2);

engine->notify_after_step_calls(/*step_calls=*/2);

References

Annotate constant arguments with a comment indicating the parameter name when calling functions or constructors. ^(link)

gemini-code-assist · 2026-05-14T08:03:20Z

+    FixedStepsScheduler scheduler(engine.get(), opt);
+
+    std::vector<std::shared_ptr<Request>> requests =
+        GenRequests({1, 1, 1}, {1, 1, 1}, RecType::kOneRec);


Constant arguments in function calls should be annotated with a comment indicating the parameter name.

Suggested change

GenRequests({1, 1, 1}, {1, 1, 1}, RecType::kOneRec);

GenRequests(/*prompt_lens=*/{1, 1, 1}, /*max_tokens=*/{1, 1, 1}, RecType::kOneRec);

References

Annotate constant arguments with a comment indicating the parameter name when calling functions or constructors. ^(link)

gemini-code-assist · 2026-05-14T08:03:20Z

@@ -3015,6 +3063,19 @@ folly::SemiFuture<std::optional<ForwardOutput>> RecWorkerImpl::step_async(
      [this, &input, index, promise = std::move(promise)]() mutable {


The input parameter is captured by reference in the asynchronous lambda. This is dangerous because step_async returns a SemiFuture, and if the caller does not block on the result immediately, the input object might be destroyed before the lambda executes, leading to a dangling reference. Capturing input by value is safer for asynchronous tasks.

Suggested change

[this, &input, index, promise = std::move(promise)]() mutable {

[this, input, index, promise = std::move(promise)]() mutable {

gemini-code-assist · 2026-05-14T08:03:20Z

+        atb::Context* atb_context = const_cast<atb::Context*>(
+            work_pipelines_[index]->runtime().context->get_atb_context());
+        atb_context->SetExecuteStream(current_stream);


The atb_context obtained from the model context should be checked for nullptr before calling SetExecuteStream. Accessing a null pointer will cause a crash. Using a CHECK assertion is recommended here to ensure the context is valid for NPU execution.

atb::Context* atb_context = const_cast<atb::Context*>(\n work_pipelines_[index]->runtime().context->get_atb_context());\n CHECK(atb_context != nullptr) << "ATB context is null for pipeline " << index;\n atb_context->SetExecuteStream(current_stream);

References

Use CHECK (glog) instead of TORCH_CHECK for assertions.

DragonFive · 2026-05-14T08:05:06Z

+                                    torch::TensorOptions()
+                                        .dtype(torch::kUInt8)
+                                        .device(beam_index.device()));
+    workspace_addr = workspace_tensor.data_ptr();


it looks like you use torch allocator for workspace to avoid malloc and free, but in multistream, it is ok for no sync but recordstream?

DragonFive · 2026-05-14T08:05:44Z

+}
+#else
+bool should_serialize_onerec_xattention_prepare(int64_t max_concurrency) {
+  UNUSED_PARAMETER(max_concurrency);


what UNUSED_PARAMETER means?

DragonFive · 2026-05-14T08:06:29Z

+      runtime_.worker.onerec_xattention_prepare_mutex_, std::defer_lock);
+  if (should_serialize_onerec_xattention_prepare(
+          runtime_.worker.options_.rec_worker_max_concurrency())) {
+    prepare_lock.lock();


do you use lock in prepare input for multi stream? may be this result in low performance?

DragonFive · 2026-05-14T08:07:58Z

+
+bool should_serialize_onerec_xattention_model_forward(int64_t max_concurrency) {
+  return max_concurrency > 1;
+}


why max_concurrency > 1 should serialize forward and prepare input, does it influence the performance?

DragonFive · 2026-05-14T08:08:48Z

+          std::defer_lock);
+      if (should_serialize_onerec_xattention_model_forward(
+              runtime_.worker.options_.rec_worker_max_concurrency())) {
+        forward_lock.lock();


why you use lock, it looks like multistream never run in same time, does gpu multistream use this too?

DragonFive · 2026-05-14T08:09:52Z

+        aclrtStream current_stream =
+            c10_npu::getCurrentNPUStream(device_.index()).stream();
+        atb::Context* atb_context = const_cast<atb::Context*>(
+            work_pipelines_[index]->runtime().context->get_atb_context());


it looks like one device only one stream? so where is multi stream?

DragonFive · 2026-05-14T08:10:49Z

why you remove it, it looks like it is used in many pipeline, will this influence other pipeline?

DragonFive added 2 commits May 13, 2026 12:03

feat: add npu onerec xattention multistream pipeline.

e757d97

feat: support onerec xattention npu scheduler multistream.

3dc61f3

DragonFive commented May 14, 2026

View reviewed changes

gemini-code-assist Bot reviewed May 14, 2026

View reviewed changes

DragonFive commented May 14, 2026

View reviewed changes

DragonFive added 4 commits May 14, 2026 16:31

bugfix: address onerec multistream review comments.

7f8a839

perf: add onerec xattention lock experiment switch.

a9ce548

perf: split onerec xattention host lock switches.

b41ffe2

perf: improve onerec xattention multistream parity.

3a05cf4

	std::unique_ptr<FakeEngine> engine = std::make_unique<FakeEngine>(64, 128);
	std::unique_ptr<FakeEngine> engine = std::make_unique<FakeEngine>(/num_blocks=/64, /block_size=/128);

	GenRequests({1, 1}, {1, 1}, RecType::kOneRec);
	GenRequests(/prompt_lens=/{1, 1}, /max_tokens=/{1, 1}, RecType::kOneRec);

	engine->notify_after_step_calls(2);
	engine->notify_after_step_calls(/step_calls=/2);

	GenRequests({1, 1, 1}, {1, 1, 1}, RecType::kOneRec);
	GenRequests(/prompt_lens=/{1, 1, 1}, /max_tokens=/{1, 1, 1}, RecType::kOneRec);

		@@ -3015,6 +3063,19 @@ folly::SemiFuture<std::optional<ForwardOutput>> RecWorkerImpl::step_async(
		[this, &input, index, promise = std::move(promise)]() mutable {

	[this, &input, index, promise = std::move(promise)]() mutable {
	[this, input, index, promise = std::move(promise)]() mutable {

Conversation

DragonFive commented May 14, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot May 14, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot May 14, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot May 14, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot May 14, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot May 14, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot May 14, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot May 14, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant