Skip to content

feat: add onerec model implement[2/N]#981

Merged
DragonFive merged 8 commits intojd-opensource:mainfrom
DragonFive:onerec/pr3-npu-full-adapt
Mar 12, 2026
Merged

feat: add onerec model implement[2/N]#981
DragonFive merged 8 commits intojd-opensource:mainfrom
DragonFive:onerec/pr3-npu-full-adapt

Conversation

@DragonFive
Copy link
Collaborator

@DragonFive DragonFive commented Mar 3, 2026

Summary

This PR is the second step of the OneRec rebaseline stack, following #970
It focuses on NPU block-layer wiring and OneRec runtime input/plumbing hardening on top of current main, while keeping existing non-
OneRec model paths stable.

What's included

  • Add OneRec block-layer contract:
    • xllm/core/layers/onerec_block_layer.h
  • Add NPU OneRec block-layer implementation:
    • xllm/core/layers/npu/npu_onerec_block_layer_impl.h
    • xllm/core/layers/npu/npu_onerec_block_layer_impl.cpp
  • Wire new OneRec layer into build targets:
    • xllm/core/layers/CMakeLists.txt
    • xllm/core/layers/npu/CMakeLists.txt
  • Extend OneRec request input handling in rec_master:
    • Validate supported tensor names (sparse_embedding, decoder_context_embedding)
    • Validate dtype/shape/content size
    • Validate required tensors and hidden-size consistency
    • Convert accepted tensors into MMData
  • Adapt runtime wiring for OneRec worker path:
    • RecWorkerImpl OneRec model init and forward path integration
    • Pipeline-side cleanup/fixes to align with current worker/executor usage
  • Scheduler guard fix:
    • Restrict REC multi-round scheduling condition to RecType::kLlmRec
  • Follow-up fixes:
    • onerec_args vocab/tokenizer argument fix
    • Device placement fix in NPU OneRec block layer (device_-based move path)
    • Build-error and merge-conflict cleanup in touched runtime code

Scope / Restrictions

  • This PR is primarily wiring and integration for OneRec on NPU path.
  • It does not introduce new non-OneRec runtime behavior.
  • The OneRec NPU block layer remains integration-level logic (no large algorithmic kernel expansion in this PR).

Backward Compatibility

  • Existing llm / vlm / embedding / non-OneRec REC flows are preserved.
  • OneRec-specific code paths are activated only for OneRec model/type selection.

Testing

  • Build-level validation for touched targets (layers, runtime, scheduler, models/rec).
  • Source-level checks for OneRec input validation branches in rec_master.
  • Smoke checks for OneRec worker forward path (encoder/decoder branch handling and device placement).
  • Regression check that non-OneRec pipelines remain unchanged.

@DragonFive DragonFive changed the title feat(rec): add onerec npu block layer wiring feat: add onerec model implement[2/N] Mar 3, 2026
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces the wiring for a new onerec_block_layer for NPU, including the abstract base class, NPU implementation, and build system changes. The implementation of NpuOneRecBlockLayerImpl::forward appears to be a placeholder, which is expected for a wiring PR. I've found a critical issue in the tensor device placement logic within the forward method that could lead to runtime errors. Please see the specific comment for details.

@DragonFive DragonFive marked this pull request as ready for review March 4, 2026 08:21
@DragonFive DragonFive force-pushed the onerec/pr3-npu-full-adapt branch from daf845a to 1bf8ce2 Compare March 7, 2026 04:20
@DragonFive DragonFive force-pushed the onerec/pr3-npu-full-adapt branch from 1bf8ce2 to 8e12383 Compare March 9, 2026 01:30
@DragonFive DragonFive merged commit 41096e2 into jd-opensource:main Mar 12, 2026
16 of 39 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants