[CPU] Refactor CPU affinity and memory management (#39781)

Signed-off-by: jiang1.li <jiang1.li@intel.com>
This commit is contained in:
Li, Jiang
2026-04-17 21:01:08 +08:00
committed by GitHub
parent b1dc87a098
commit d02421a7db
14 changed files with 721 additions and 423 deletions
+4
View File
@@ -141,6 +141,8 @@ void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc,
torch::Tensor slot_mapping,
const int64_t block_size);
void init_cpu_memory_env(std::vector<int64_t> node_ids);
namespace cpu_utils {
void eagle_prepare_inputs_padded_kernel_impl(
const torch::Tensor& cu_num_draft_tokens,
@@ -431,6 +433,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
"block_size) -> ()",
&compute_slot_mapping_kernel_impl);
ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env);
// Speculative decoding kernels
ops.def(
"eagle_prepare_inputs_padded_kernel_impl(Tensor cu_num_draft_tokens, "