Align Operators

Align Operators

Use this skill for operator-level precision alignment only.

Core Rule

Treat each operator as a mathematical function and compare outputs for the same synthetic input.

Do not use this skill for full transformer runtime, batch semantics, sequence semantics, KV cache behavior, or distributed execution.

Fused Operators Note

Rust may implement fused operators (e.g., fused SiLU + multiply, fused add + RMS norm) for performance optimization, while HuggingFace/Python will implement them as separate operators. When aligning:

• Do not worry about implementation details (fused vs non-fused)
• Only focus on input/output mathematical equivalence
• For a fused Rust operator, generate the expected output by running the individual HF operators sequentially on the same input

Standard Setup

• Use a realistic model configuration.
• Prefer hidden_size=2048, intermediate_size=5472, moe_intermediate_size=768, num_attention_heads=32, num_key_value_heads=4, head_dim=128, num_experts=128, num_experts_per_tok=8.
• Keep tensors lightweight but representative, such as [4, 2048], [4, 32, 128], and [4, 4, 128].
• Use strict FP32 in both Rust and Python.
• Keep execution deterministic.

Recommended Workflow

1. Generate synthetic inputs in Python and save them as input.npy.
2. Run the Rust operator (fused or not) on the same input and save output.npy.
3. For a fused Rust operator, run the equivalent sequence of individual HuggingFace/PyTorch operators in FP32 to generate the reference output.
4. Compare shapes, max absolute error, mean absolute error, and cosine similarity.
5. Accept only when the outputs meet the FP32 thresholds.

Input Strategy

Use synthetic inputs only:

• sequential values for layout and reshape checks
• all zeros for stability and epsilon checks
• all ones for scaling checks
• small random values for general correctness
• extreme values for softmax, SiLU, and rsqrt edge cases

Target Operators

Focus on:

• Linear
• RMSNorm
• SiLU
• SwiGLU
• Softmax
• Attention
• Fused combinations (e.g., SiLU+Multiply, Add+RMSNorm, MatMul+Sigmoid)

Comparison Rules

• Use NumPy .npy for interchange.
• Compare the Rust output against the HuggingFace output for the same mathematical operation (even if implementation differs).
• For fused Rust operators, compare against sequential HF operators.
• Flag layout, reshape, transpose, softmax stability, and implicit f64 issues.

Project Structure

Each operator (including fused ones) has its own directory under alignment/:

alignment/
├── README.md
├── rope/
│   ├── dump/              # Test outputs (npy, bin files)
│   ├── generate_hf_rope.py      # Python reference
│   └── test_rope_alignment.py   # Comparison script
├── silu_mul/              # Fused SiLU + multiply example
│   ├── dump/
│   ├── generate_hf_silu_mul.py
│   └── test_silu_mul_alignment.py
├── linear/
├── rmsnorm/
└── [operator_name]/

To add a new operator, create a directory following the rope/ pattern.

Environment Tips

Windows-Specific Issues

1. Python Command: On Windows, use py instead of python to run Python scripts:

   py alignment/[operator]/generate_hf_[operator].py
   

   This avoids environment/path issues that may occur with python command.

2. Rust Build Cache: If you encounter build errors like "failed to run custom build command" for dependencies like serde or quote:

   # Clean the target build directory
   rm -r target
   # Then rebuild
   cargo build
   

3. PyTorch Installation: Ensure PyTorch is installed in the environment used by the py command. Use py -m pip to install:

   py -m pip install torch numpy
   

Reference Material

See references/alignment-spec.md for the full alignment spec, tensor-shape matrix, thresholds, logging format, bug checklist, and detailed guide for adding new operators.