Skip to content

Conversation

@emharsha1812
Copy link
Contributor

@emharsha1812 emharsha1812 commented Dec 29, 2025

Description

This PR adds a new Medium difficulty challenge: Rotary Positional Embedding (RoPE).

RoPE is a critical component in modern LLMs (Llama 2/3, Mistral, etc.). I noticed that Unsloth recently added optimized RoPE kernels to their main repo, so I thought it would be a great addition to LeetGPU for users to practice implementing this fundamental operation.

Challenge Details

  • Path: challenges/medium/61_rope_embedding
  • Concept: Implementing the rotation formula efficiently on GPU.
  • Reference Implementation: PyTorch-based reference in challenge.py.
  • Starters: Added templates for CUDA, Triton, Mojo, PyTorch, and CuTe.

Verification

  • challenge.py reference implementation verified against standard formula.

@kunal-mansukhani
Copy link
Contributor

@23silicon Can you make the spec the correct format, add more functional tests, and add JAX starter code

import jax
import jax.numpy as jnp

#Q, cos, sin are tensors on the GPU
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Need space between # and Q


#Q, cos, sin are tensors on the GPU
@jax.jit
def solve(Q: array, cos: array, sin: array, output: array, M: int, D: int) -> jax.Array:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

jax.Array, not array

Copy link
Contributor

@kunal-mansukhani kunal-mansukhani Jan 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

no output. It's returned

@kunal-mansukhani kunal-mansukhani self-requested a review January 2, 2026 01:41
@kunal-mansukhani kunal-mansukhani merged commit 4e72dfb into AlphaGPU:main Jan 2, 2026
4 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants