feat: add xAttention support for Qwen3 generative recommendation. #586

LMX-xin · 2025-12-23T08:34:37Z

Adapted xAttention for the Qwen3 model in generative recommendation scenarios.
- Added xAttention-compatible scheduling logic that merges multiple steps into a single step, enabling the full workflow from prefill to multiple decode iterations to run entirely on the device side.
- Added xAttention-compatible KV cache management logic.
- When xAttention is enabled, also integrated and adapted the related ACL Graph workflow.

xllm/core/distributed_runtime/llm_engine.cpp

xllm/core/distributed_runtime/llm_master.cpp

XuZhang99 · 2025-12-23T08:58:41Z

xllm/core/framework/batch/multi_step_batch_input_builder.h

I think you should inherit BatchInputBuilder not create a new standalone class.

XuZhang99 · 2025-12-23T09:00:03Z

xllm/core/framework/request/sequences_group.cpp

+    std::vector<SequenceOutput>& outputs,
+    const Tokenizer& tokenizer,
+    const Sequence& base) {
+  VLOG(1) << "[debug] generate_multi_round_output, bw: "


what is VLOG(1)? why not use LOG?

XuZhang99 · 2025-12-23T09:02:08Z

xllm/core/kernels/npu/xllm_ops/beam_search_group.h

+#include "acl/acl.h"
+#include "aclnn_beam_search_group.h"
+#include "acltensor_utils.h"
+#include "util/tensor_helper.h"


there are many repetitive headers included in .h and .cpp file in xllm_ops dir.

XuZhang99 · 2025-12-23T09:03:08Z

xllm/core/kernels/npu/xllm_ops/cache_select.cpp

+#include "aclnn_select_unshared_kv.h"
+#include "cache_select.h"
+
+#define CHECK_ACL_SUCCESS(expr, msg) \


move this macro to acltensor_utils.h.

xllm/core/runtime/llm_worker_impl.cpp

XuZhang99 · 2025-12-23T09:33:10Z

xllm/core/runtime/llm_worker_impl.cpp

+#include <tuple>
+
+#include "kernels/npu/xllm_ops/beam_search_group.h"
+#include "kernels/npu/xllm_ops/cache_select.h"


you need add xllm_ops_api.h, like mlu_ops_api.h and cuda_ops_api.h.

xllm/core/distributed_runtime/llm_engine.cpp

xllm/core/framework/batch/multi_step_batch_input_builder.cpp

xllm/core/runtime/params_utils.cpp

xllm/core/runtime/acl_graph_executor_rec_impl.cpp

xllm/core/runtime/worker_impl.cpp

xllm/proto/worker.proto

…into factory.

DragonFive · 2026-01-09T09:22:54Z

/gemini review

gemini-code-assist

Code Review

This pull request introduces support for xAttention in Qwen3 models for generative recommendation scenarios. The changes are extensive, adding new scheduling logic, KV cache management, and adapting the ACL Graph workflow for multi-step decoding. The introduction of QwenRecWorkerImpl and AclGraphRecExecutorImpl effectively isolates the new functionality. However, I have identified two significant issues. There's a high-severity issue concerning potentially excessive memory allocation for the shared KV cache, and a critical issue where a KV cache capacity check is bypassed, which could lead to memory corruption. These should be addressed to ensure the stability and efficiency of the new feature.

gemini-code-assist · 2026-01-09T09:30:20Z

xllm/core/framework/request/sequence_kv_state.cpp

 void KVCacheState::incr_kv_cache_tokens_num(size_t num) {
+  if (FLAGS_max_decode_rounds > 0) {
+    kv_cache_tokens_num_ += num;
+    return;
+  }
  CHECK(kv_cache_tokens_num_ + num <= current_max_tokens_capacity());
  kv_cache_tokens_num_ += num;
 }


Bypassing the KV cache capacity check when FLAGS_max_decode_rounds > 0 is dangerous. This removes a critical safety measure that prevents out-of-bounds writes to the KV cache. Even if a different memory management scheme is used for multi-round decoding, there is still a capacity limit for any allocated buffer. Without this check, it's possible for kv_cache_tokens_num_ to exceed the allocated capacity, which could lead to memory corruption and unpredictable behavior. This check should be restored or replaced with an equivalent check appropriate for the new memory scheme.

void KVCacheState::incr_kv_cache_tokens_num(size_t num) { CHECK(kv_cache_tokens_num_ + num <= current_max_tokens_capacity()); kv_cache_tokens_num_ += num; }

gemini-code-assist · 2026-01-09T09:30:20Z

xllm/core/framework/batch/multi_step_batch_input_builder.cpp

+    raw_forward_input.shared_kv_shape = {
+        batch_size * FLAGS_max_token_per_req, n_kv_heads, head_dim};


The calculation for shared_kv_shape uses FLAGS_max_token_per_req as a multiplier for the batch size. This can lead to significant over-allocation of memory for the shared KV cache, especially if FLAGS_max_token_per_req is set to a large value to accommodate a few long requests, while the average request length is much shorter. This could result in unnecessary memory pressure and potential out-of-memory errors. As noted in the TODO comment on lines 339-342, it would be much safer and more memory-efficient to calculate the required size based on the actual token needs of the sequences in the current batch.

LMX-xin requested review from DragonFive, RobbieLeung, liutongxuan and yq33victor December 23, 2025 08:34

XuZhang99 reviewed Dec 23, 2025

View reviewed changes