[Bug] Incorrect RoPE implementation (adjacent pairing vs GPT-J / HF style) causes Output is incorrect.

[Rlingb]

Describe the bug

When running inference with Llama-3.2-1B-Instruct, the model produces abnormal outputs.
For example(after detokenization):

1. Prompt: "4+7=?" → Output: "+"
2. Prompt: "My name is Ash." → Output: "How's name is Ash"

Root Cause

The issue lies in the RoPE (Rotary Position Embedding) implementation in kernels.cu.
The current code uses adjacent pairing for rotary dimensions:

// Current (WRONG) implementation in ropeKernel:
__nv_bfloat16 prev_2i = input[2 * threadIdx.x + blockIdx.x * proj_dim];
__nv_bfloat16 prev_2i_1 = input[2 * threadIdx.x + 1 + blockIdx.x * proj_dim];
However, Llama 3 (and Hugging Face transformers) uses GPT-J style pairing, where dimensions are paired as (i, i + head_dim/2) instead of (2i, 2i+1).

Since the pretrained weights of Llama-3.2-1B were optimized under GPT-J pairing, the current adjacent pairing misaligns the rotational angles with the learned feature dimensions, causing the Q/K dot products to lose positional semantics entirely.

Proposed Fix

In my local fork, I've been experimenting with a potential fix. Beyond adjusting the pairing strategy to GPT-J style, I also tried precomputing the cos/sin tables to reduce computational overhead, and I implemented the Llama 3 three-band frequency scaling。
I've attached a reference implementation of the kernel below for your convenience (please feel free to point out if anything looks off):

__global__ void ropeKernel_llama3(__nv_bfloat16 *input, int num_tokens, int proj_dim,
                                  int head_dim, const float *cos_table, const float *sin_table)
{
    int token_idx = blockIdx.x;
    int tid = threadIdx.x;
    int half_proj = proj_dim / 2;
    int half_dim = head_dim / 2;

    if (tid >= half_proj)
        return;

    int head_idx = tid / half_dim;
    int pair_idx = tid % half_dim;

    int base = token_idx * proj_dim + head_idx * head_dim;
    int idx1 = base + pair_idx;
    int idx2 = base + pair_idx + half_dim;

    float x1 = (float)input[idx1];
    float x2 = (float)input[idx2];

    int table_idx = token_idx * head_dim + pair_idx * 2;
    float c = cos_table[table_idx];
    float s = sin_table[table_idx];

    input[idx1] = (__nv_bfloat16)(x1 * c - x2 * s);
    input[idx2] = (__nv_bfloat16)(x1 * s + x2 * c);
}

After applying this fix, I tested the same prompts:

1. "4+7=?" → “4 + 7 = 11”
2. "My name is Ash." → “Nice to meet you, Ash! Is there something I can help with or would you like to chat?”

In the end

I’m keen to contribute these changes back to the project, and I’ve finished writing up a full PR in my local branch. Rather than opening it right away, I was hoping to get your feedback first, so I can align the implementation properly with your vision for tiny-vllm.

Could you let me know if you’re happy to review and merge this PR? Any suggestions for revisions are totally welcome.

[jmaczan]

Hi, thx for letting me know! Does this PR address the issue? #1 If yes, can you help with verification and testing?

[Rlingb]

@jmaczan Yes, this PR addresses the issue. I've already submitted it – please take a look when you have time. And feel free to let me know if you need any help with testing or further adjustments. Thanks! 🙏

[jmaczan]

I need some quick way of testing it without spending too much time on it, please show tests or something that will show clearly that everything works as expected

[Rlingb]

@jmaczan Sure! Here's a quick and minimal test using the built-in prompt in main.cpp.

Currently, the test prompt is:

// PROMPT 0 (What is 2+2?) - length 17
std::queue<std::vector<int>> queue;
queue.push({128000, 128006, 882, 128007, 271, 3923, 374, 220, 17, 10, 17, 30, 128009, 128006, 78191, 128007, 271});

You can replace it with the tokenized version of "4+7=?" to verify the fix:

queue.push({128000, 128006, 882, 128007, 271, 19, 10, 22, 20477, 128009, 128006, 78191, 128007, 271});
Expected result:

Before fix (buggy): + (or gibberish)

After fix (this PR): 4 + 7 = 11

I've already tested this locally and confirmed the fix works. Let me know if you need any more details!

[Rlingb]

@jmaczan Is the test going well?