Skip to content

Commit 64e3ba3

Browse files
[CPU] Add native LTX-Video RoPE kernel and enable the fusion (#37141)
### Details LTX-Video's 3-axis RoPE was left decomposed on CPU because the fusion was disabled, the existing interleaved executor can't handle LTX's 3D layout with separate full-width cos/sin tables. This PR adds a native RoPEExecutorLtxVideo that computes the rotation in f32 and rounds once to the output type, enables RoPEFusionLtxVideo on CPU (same fused op as GPU), and instantiates the shared RoPETestLtxVideo test for CPU. RoPE subgraph, seq 2520 × 2048, median CPU latency, fused vs decomposed: | precision | decomposed | fused | speedup | |-----------|------------|--------|---------| | f32 | 23.05 ms | 11.26 ms | **2.05×** | | bf16 | 23.19 ms | 14.69 ms | **1.58×** | f16 stays on the decomposed path for now. unlike bf16, its pipeline runs ConvertPrecision before RoPEFusion, so the fusion doesn't fire. Enabling f16 fusion is a follow-up. and is anyways producing NaN as of now (ref #37039). ```mermaid flowchart LR subgraph Before["Before — decomposed (~8 ops)"] direction LR x1[x] --> R1[Reshape] R1 --> SP[Split] SP -->|imag| NEG["Multiply(-1)"] NEG --> CC[Concat] SP -->|real| CC CC --> R2[Reshape] x1 --> MC["Multiply · cos"] cos1[cos] --> MC R2 --> MS["Multiply · sin"] sin1[sin] --> MS MC --> ADD[Add] MS --> ADD ADD --> o1[out] end subgraph After["After — fused"] direction LR x2[x] --> ROPE["RoPE(is_ltx_video)"] cos2[cos] --> ROPE sin2[sin] --> ROPE ROPE --> o2[out] end ``` ### AI Assistance: - *AI assistance used: no / yes* yes - *If yes, summarize how AI was used and what human validation was performed (build/tests/manual checks).*
1 parent be0a889 commit 64e3ba3

4 files changed

Lines changed: 59 additions & 1 deletion

File tree

src/plugins/intel_cpu/src/nodes/rope.cpp

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -244,6 +244,46 @@ struct RoPE::RoPEExecutorInterleaved : public RoPE::Executor {
244244
}
245245
};
246246

247+
// LTX-Video 3D spatial-temporal RoPE: x [batch, seq, rotary_ndims] with separate full-width cos/sin
248+
// tables. Interleaved complex pairs; each element keeps its own cos/sin (the two halves of a pair
249+
// need not share an angle). Accumulated in f32 (cos/sin ports are f32) and rounded once on store,
250+
// so bf16 stays as precise as PyTorch.
251+
template <typename T>
252+
struct RoPE::RoPEExecutorLtxVideo : public RoPE::Executor {
253+
const op::internal::RoPE::Config& m_config;
254+
255+
explicit RoPEExecutorLtxVideo(const op::internal::RoPE::Config& config) : m_config(config) {}
256+
257+
void execute([[maybe_unused]] const dnnl::stream& strm,
258+
const std::vector<MemoryPtr>& inputs,
259+
const std::vector<MemoryPtr>& outputs,
260+
const CpuParallelPtr& cpu_parallel) override {
261+
ov::intel_cpu::PlainTensor t_src(inputs[0]);
262+
ov::intel_cpu::PlainTensor t_cos(inputs[1]);
263+
ov::intel_cpu::PlainTensor t_sin(inputs[2]);
264+
ov::intel_cpu::PlainTensor t_dst(outputs[0]);
265+
266+
auto batch_size = t_src.size(0);
267+
auto seq_len = t_src.size(1);
268+
auto rotary_dims = m_config.rotary_ndims;
269+
270+
cpu_parallel->parallel_for2d(batch_size, seq_len, [&](size_t b, size_t p) {
271+
auto* x = t_src.ptr<T>(b, p);
272+
// allow_broadcast handles size-1 cos/sin batch/seq natively
273+
const float* cos = &t_cos.at<float>({b, p, 0}, true);
274+
const float* sin = &t_sin.at<float>({b, p, 0}, true);
275+
auto* dst = t_dst.ptr<T>(b, p);
276+
277+
for (size_t r = 0; r < rotary_dims; r += 2) {
278+
auto real = static_cast<float>(x[r]);
279+
auto imag = static_cast<float>(x[r + 1]);
280+
dst[r] = static_cast<T>(cos[r] * real - sin[r] * imag);
281+
dst[r + 1] = static_cast<T>(sin[r + 1] * real + cos[r + 1] * imag);
282+
}
283+
});
284+
}
285+
};
286+
247287
template <typename T>
248288
struct RoPE::RoPEExecutorChatGLM : public RoPE::Executor {
249289
const op::internal::RoPE::Config& m_config;
@@ -470,6 +510,17 @@ void RoPE::initSupportedPrimitiveDescriptors() {
470510
m_executor = std::make_shared<RoPEExecutorChatGLM<float>>(m_config);
471511
rtPrecision = ov::element::f32;
472512
}
513+
} else if (m_config.is_ltx_video) {
514+
CPU_NODE_ASSERT(m_config.rotary_ndims % 2 == 0, "rotary_ndims must be even for LTX RoPE");
515+
// LTX sets both is_interleaved and is_ltx_video, so this must be checked first
516+
if (rtPrecision == ov::element::f16) {
517+
m_executor = std::make_shared<RoPEExecutorLtxVideo<ov::float16>>(m_config);
518+
} else if (rtPrecision == ov::element::bf16) {
519+
m_executor = std::make_shared<RoPEExecutorLtxVideo<ov::bfloat16>>(m_config);
520+
} else {
521+
m_executor = std::make_shared<RoPEExecutorLtxVideo<float>>(m_config);
522+
rtPrecision = ov::element::f32;
523+
}
473524
} else if (m_config.is_interleaved) {
474525
CPU_NODE_ASSERT(m_config.slice_start == 0, "slice_start must be 0 for interleaved mode");
475526
CPU_NODE_ASSERT(m_config.slice_stop == 0, "slice_stop must be 0 for interleaved mode");

src/plugins/intel_cpu/src/nodes/rope.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,8 @@ class RoPE : public Node {
4949
template <typename T>
5050
struct RoPEExecutorInterleaved;
5151
template <typename T>
52+
struct RoPEExecutorLtxVideo;
53+
template <typename T>
5254
struct RoPEExecutorChatGLM;
5355
template <typename T>
5456
struct RoPEExecutorQwen;

src/plugins/intel_cpu/src/transformations/transformation_pipeline.cpp

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1139,7 +1139,6 @@ void Transformations::PostLpt() {
11391139
CPU_REGISTER_PASS_X64(postLPTPassManager, ov::pass::RoPEFusion, true);
11401140
CPU_REGISTER_PASS_ARM64(postLPTPassManager, ov::pass::RoPEFusion, true);
11411141
CPU_DISABLE_PASS_COMMON(postLPTPassManager, ov::pass::RoPEFusionFlux);
1142-
CPU_DISABLE_PASS_COMMON(postLPTPassManager, ov::pass::RoPEFusionLtxVideo);
11431142
CPU_REGISTER_PASS_X64(postLPTPassManager, CausalMaskPreprocessFusion);
11441143

11451144
#if defined(OPENVINO_ARCH_X86_64)

src/plugins/intel_cpu/tests/functional/shared_tests_instances/subgraph_tests/rotary_pos_emb.cpp

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -93,5 +93,11 @@ INSTANTIATE_TEST_SUITE_P(smoke_RoPETestGPTOSS,
9393
::testing::Values(ov::test::utils::DEVICE_CPU)),
9494
RoPETestGPTOSS::getTestCaseName);
9595

96+
INSTANTIATE_TEST_SUITE_P(smoke_RoPETestLtxVideo,
97+
RoPETestLtxVideo,
98+
::testing::Combine(::testing::Values(ov::element::f32),
99+
::testing::Values(ov::test::utils::DEVICE_CPU)),
100+
RoPETestLtxVideo::getTestCaseName);
101+
96102
} // namespace test
97103
} // namespace ov

0 commit comments

Comments
 (0)