Llama-lineage CPU operations — RmsNorm, Swiglu, Rope, TokenEmbedding

[ToddThomson]

Area: Mila/Src/Dnn/Compute/Devices/Cpu/Operations/

Summary

Mila's CPU backend implements the GPT-2 operation set (Linear, Gelu, LayerNorm,
Residual, Softmax, MHA, LPE, Sampling) but not the Llama/Gemma-lineage ops.
As a result, Llama and Gemma models cannot run on the CPU device — dispatch fails
at compile time with a missing OperationTraits<...> specialization (a hard error
by design, not a runtime miss). This issue adds the four net-new CPU forward ops
so the Llama-lineage forward pass has a CPU path.

This is demand-driven, not a release gate: absence is zero-cost on the GPU
path (OperationTraits.Cpu.ixx header notes CPU is FP32-only and full CPU parity
is not required). It's a well-bounded task with clear reference implementations,
which is what makes it a good first issue.

Scope

Add FP32-only concrete CPU ops for the four Llama-lineage operations that have no
CPU implementation today:

Op	New file	CUDA reference
RmsNormOp	CpuRmsNormOp.ixx	Cuda/Operations/.../CudaRmsNormOp*
SwigluOp	CpuSwigluOp.ixx	Cuda/Operations/Activations/Swiglu/CudaSwigluOp.ixx
RopeOp	CpuRopeOp.ixx	Cuda/Operations/Encodings/Rope/CudaRopeOp.ixx
TokenEmbeddingOp	CpuTokenEmbeddingOp.ixx	Cuda/Operations/Embeddings/CudaTokenEmbeddingOp.ixx

Out of scope (deliberately): CrossEntropyOp — a CpuCrossEntropyOp.ixx /
CpuSoftmaxCrossEntropyOp.ixx already exist on disk but are unwired
(OperationTraits.Cpu.ixx:19 marks it "pending"), and wiring the loss op is
coupled to the separate loss/backward-path revival. Leave it for that work.

Pattern to follow

Each op is a concrete (non-templated) FP32 class mirroring the existing
CpuGeluOp.ixx / CpuLinearOp.ixx / CpuLayerNormOp.ixx. For each of the four:

1. Write CpuXxxOp.ixx — a straightforward FP32 host implementation of the
   forward math (port the numerics from the CUDA reference kernel; backward is not
   required for these inference-path ops, matching the other CPU ops).
2. Register it in
   OperationTraits.Cpu.ixx:
   add the import Compute.CpuXxxOp; and a specialization, e.g.:

   template<>
   struct OperationTraits<OperationType::RmsNormOp, DeviceType::Cpu, TensorDataType::FP32, void>
   {
       using type = CpuRmsNormOp;
   };

   (LinearOp uses NoWeightQuant as its policy slot; the others use void —
   follow the existing entries in the file.)
3. Wire the new .ixx into the CPU operations CMake target (alongside the other
   Cpu*Op.ixx) and update the module aggregator CpuOperations.ixx if it
   re-exports the op modules.
4. Update the "Migration status" table in the OperationTraits.Cpu.ixx file
   header (mark each newly-added op complete).

Acceptance criteria

• Each new op has an OperationTraits<Op, Cpu, FP32, ...> specialization and
  compiles into the CPU build.
• A Llama<Cpu, FP32> (or the component-level RmsNorm/Swiglu/Rope/TokenEmbedding
  on the CPU device) constructs and runs its forward pass without a missing-
  specialization compile error.
• Per-op CPU unit tests (mirror the existing Cpu-device component tests):
  each op's FP32 forward output matches a hand-computed or CUDA-cross-checked
  reference within tolerance. The tests run under the CPU-only CI ratchet
  (cpu-only-tests), so they guard against regression automatically.
• OperationTraits.Cpu.ixx migration-status table updated.

References

• Existing CPU op pattern: Cpu/Operations/CpuGeluOp.ixx, CpuLinearOp.ixx,
  CpuLayerNormOp.ixx
• Dispatch design: Mila/Specifications/OperationDispatch.md
• CUDA reference numerics: the Cuda*Op.ixx files in the table above