-
Add Metal rotary embedding kernel matching vLLM interface
New Metal kernel implementing rotary_embedding(positions, query, key,
head_size, cos_sin_cache, is_neox) with the exact signature vLLM expects.
Features:
- Supports fp16, bf16, fp32 dtypes
- NeoX style (Llama, Mistral) and GPT-J style rotation
- Arbitrary head dims (64, 128, 256)
- GQA support (separate num_heads / num_kv_heads)
- Optional key tensor (key=None skips key rotation)
- Function constants for IS_NEOX (zero-cost specialization)
- Precomputed cos_sin_cache lookup (not on-the-fly frequency computation)
Includes comprehensive tests against pure-PyTorch reference implementation.
Co-developed-by: Claude Code v2.1.50 (claude-opus-4-6)
