[BugFix] Accuracy fix for llama4 int4 - improperly casted scales

[LucasWilkinson]

A few MoE fixes:

1. Scales where fp16/bf16 but were being casted to float unsafely. Add a check by using .data_ptr<float>()
2. Fix AttributeError: 'FusedMoE' object has no attribute 'params_dtype'
3. Make sure Llama4 returns float router weights to ensure compatibility with kernels (namely moe_wna16_gemm )

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[liuzijing2014]

If top_k = 2, then for each token, we would have 2 top_k_weights to be applied. Looking at the line here, it assumes topk_weights[token_index] is always a scalar value?

[LucasWilkinson]

its a bit weird but the way it works is that token_index ranges between 0 and topk * num_tokens, then topk_weights is of size topk * num_tokens, and basically we index the input via input[token_index / topk] but the topk_weights[token_index]

e.g. if we had the router output as topk 2 with 2 tokens:

weights = [
  [0.1, 0.2], 
  [0.3, 0.4]
]
topk_ids = [
  [1, 0],
  [1, 2]
]

and if we had block size 1 for MoE align block size then we'd get:

sorted_token_ids = [1, 0, 2, 3]
expert_ids = [0, 1, 1, 2]

so token_index (a poor name) is actually the index from sorted_token_ids which is of size topk * num_tokens, meaning topk_weights[token_index] is still a scalar

[liuzijing2014]

I see so basically for topk=2, the outer loop of for (int m = 0; m < num_valid_tokens; ++m) would iterate each token in the request twice, the first time it get token x first expert weight, and the second time, it get token x second expert weight. Is this the right understanding?

[liuzijing2014]

Confirmed the fix recover llama4 int4 checkpoint accuracy back to normal.

[mgoin]

Thanks so much for this bundle of fixes, LGTM! cc @jinzhen-lin for the moe_wna16 cuda update

[liuzijing2014]

INT4 Scout

TP=8 after the change

|       Tasks       |Version|    Filter    |n-shot|  Metric   |   |Value |   |Stderr|
|-------------------|------:|--------------|-----:|-----------|---|-----:|---|-----:|
|mmlu_pro           |    2.0|custom-extract|      |exact_match|↑  |0.7168|±  |0.0040|
| - biology         |    2.1|custom-extract|     5|exact_match|↑  |0.8201|±  |0.0144|
| - business        |    2.1|custom-extract|     5|exact_match|↑  |0.7592|±  |0.0152|
| - chemistry       |    2.1|custom-extract|     5|exact_match|↑  |0.7739|±  |0.0124|
| - computer_science|    2.1|custom-extract|     5|exact_match|↑  |0.6829|±  |0.0230|
| - economics       |    2.1|custom-extract|     5|exact_match|↑  |0.8081|±  |0.0136|
| - engineering     |    2.1|custom-extract|     5|exact_match|↑  |0.6161|±  |0.0156|
| - health          |    2.1|custom-extract|     5|exact_match|↑  |0.7311|±  |0.0155|
| - history         |    2.1|custom-extract|     5|exact_match|↑  |0.6037|±  |0.0251|
| - law             |    2.1|custom-extract|     5|exact_match|↑  |0.4732|±  |0.0151|
| - math            |    2.1|custom-extract|     5|exact_match|↑  |0.8157|±  |0.0106|
| - other           |    2.1|custom-extract|     5|exact_match|↑  |0.6732|±  |0.0154|
| - philosophy      |    2.1|custom-extract|     5|exact_match|↑  |0.5832|±  |0.0221|
| - physics         |    2.1|custom-extract|     5|exact_match|↑  |0.7898|±  |0.0113|
| - psychology      |    2.1|custom-extract|     5|exact_match|↑  |0.7682|±  |0.0149|

| Tasks |Version|Filter|n-shot|     Metric      |   |Value |   |Stderr|
|-------|------:|------|-----:|-----------------|---|-----:|---|-----:|
|chartqa|      0|none  |     0|anywhere_accuracy|↑  |0.8900|±  |0.0063|
|       |       |none  |     0|exact_match      |↑  |0.6628|±  |0.0095|
|       |       |none  |     0|relaxed_accuracy |↑  |0.8884|±  |0.0063|

TP=1 after the change

|       Tasks       |Version|    Filter    |n-shot|  Metric   |   |Value |   |Stderr|
|-------------------|------:|--------------|-----:|-----------|---|-----:|---|-----:|
|mmlu_pro           |    2.0|custom-extract|      |exact_match|↑  |0.7186|±  |0.0040|
| - biology         |    2.1|custom-extract|     5|exact_match|↑  |0.8312|±  |0.0140|
| - business        |    2.1|custom-extract|     5|exact_match|↑  |0.7541|±  |0.0153|
| - chemistry       |    2.1|custom-extract|     5|exact_match|↑  |0.7739|±  |0.0124|
| - computer_science|    2.1|custom-extract|     5|exact_match|↑  |0.7024|±  |0.0226|
| - economics       |    2.1|custom-extract|     5|exact_match|↑  |0.8057|±  |0.0136|
| - engineering     |    2.1|custom-extract|     5|exact_match|↑  |0.6120|±  |0.0157|
| - health          |    2.1|custom-extract|     5|exact_match|↑  |0.7408|±  |0.0153|
| - history         |    2.1|custom-extract|     5|exact_match|↑  |0.6037|±  |0.0251|
| - law             |    2.1|custom-extract|     5|exact_match|↑  |0.4868|±  |0.0151|
| - math            |    2.1|custom-extract|     5|exact_match|↑  |0.8075|±  |0.0107|
| - other           |    2.1|custom-extract|     5|exact_match|↑  |0.6764|±  |0.0154|
| - philosophy      |    2.1|custom-extract|     5|exact_match|↑  |0.6032|±  |0.0219|
| - physics         |    2.1|custom-extract|     5|exact_match|↑  |0.7860|±  |0.0114|
| - psychology      |    2.1|custom-extract|     5|exact_match|↑  |0.7619|±  |0.0151|

| Tasks |Version|Filter|n-shot|     Metric      |   |Value |   |Stderr|
|-------|------:|------|-----:|-----------------|---|-----:|---|-----:|
|chartqa|      0|none  |     0|anywhere_accuracy|↑  |0.8868|±  |0.0063|
|       |       |none  |     0|exact_match      |↑  |0.6588|±  |0.0095|
|       |       |none  |     0|relaxed_accuracy |↑  |0.8844|±  |0.0064|

Before the change

|       Tasks       |Version|    Filter    |n-shot|  Metric   |   |Value |   |Stderr|
|-------------------|------:|--------------|-----:|-----------|---|-----:|---|-----:|
|mmlu_pro           |    2.0|custom-extract|      |exact_match|↑  |0.0286|±  |0.0013|
| - biology         |    2.1|custom-extract|     5|exact_match|↑  |0.0126|±  |0.0042|
| - business        |    2.1|custom-extract|     5|exact_match|↑  |0.0000|±  |0.0000|
| - chemistry       |    2.1|custom-extract|     5|exact_match|↑  |0.0168|±  |0.0038|
| - computer_science|    2.1|custom-extract|     5|exact_match|↑  |0.0195|±  |0.0068|
| - economics       |    2.1|custom-extract|     5|exact_match|↑  |0.0012|±  |0.0012|
| - engineering     |    2.1|custom-extract|     5|exact_match|↑  |0.0093|±  |0.0031|
| - health          |    2.1|custom-extract|     5|exact_match|↑  |0.0000|±  |0.0000|
| - history         |    2.1|custom-extract|     5|exact_match|↑  |0.4777|±  |0.0256|
| - law             |    2.1|custom-extract|     5|exact_match|↑  |0.0091|±  |0.0029|
| - math            |    2.1|custom-extract|     5|exact_match|↑  |0.0785|±  |0.0073|
| - other           |    2.1|custom-extract|     5|exact_match|↑  |0.0000|±  |0.0000|
| - philosophy      |    2.1|custom-extract|     5|exact_match|↑  |0.0000|±  |0.0000|
| - physics         |    2.1|custom-extract|     5|exact_match|↑  |0.0000|±  |0.0000|
| - psychology      |    2.1|custom-extract|     5|exact_match|↑  |0.0000|±  |0.0000|

BF16 Scout TP=8

After the change

|       Tasks       |Version|    Filter    |n-shot|  Metric   |   |Value |   |Stderr|
|-------------------|------:|--------------|-----:|-----------|---|-----:|---|-----:|
|mmlu_pro           |    2.0|custom-extract|      |exact_match|↑  |0.7090|±  |0.0040|
| - biology         |    2.1|custom-extract|     5|exact_match|↑  |0.6764|±  |0.0175|
| - business        |    2.1|custom-extract|     5|exact_match|↑  |0.7617|±  |0.0152|
| - chemistry       |    2.1|custom-extract|     5|exact_match|↑  |0.7747|±  |0.0124|
| - computer_science|    2.1|custom-extract|     5|exact_match|↑  |0.6951|±  |0.0228|
| - economics       |    2.1|custom-extract|     5|exact_match|↑  |0.8128|±  |0.0134|
| - engineering     |    2.1|custom-extract|     5|exact_match|↑  |0.6151|±  |0.0156|
| - health          |    2.1|custom-extract|     5|exact_match|↑  |0.7262|±  |0.0156|
| - history         |    2.1|custom-extract|     5|exact_match|↑  |0.5932|±  |0.0252|
| - law             |    2.1|custom-extract|     5|exact_match|↑  |0.4696|±  |0.0150|
| - math            |    2.1|custom-extract|     5|exact_match|↑  |0.8120|±  |0.0106|
| - other           |    2.1|custom-extract|     5|exact_match|↑  |0.6851|±  |0.0153|
| - philosophy      |    2.1|custom-extract|     5|exact_match|↑  |0.5912|±  |0.0220|
| - physics         |    2.1|custom-extract|     5|exact_match|↑  |0.7860|±  |0.0114|
| - psychology      |    2.1|custom-extract|     5|exact_match|↑  |0.7744|±  |0.0148|

| Tasks |Version|Filter|n-shot|     Metric      |   |Value |   |Stderr|
|-------|------:|------|-----:|-----------------|---|-----:|---|-----:|
|chartqa|      0|none  |     0|anywhere_accuracy|↑  |0.8892|±  |0.0063|
|       |       |none  |     0|exact_match      |↑  |0.6528|±  |0.0095|
|       |       |none  |     0|relaxed_accuracy |↑  |0.8856|±  |0.0064|

Before the change

|       Tasks       |Version|    Filter    |n-shot|  Metric   |   |Value |   |Stderr|
|-------------------|------:|--------------|-----:|-----------|---|-----:|---|-----:|
|mmlu_pro           |    2.0|custom-extract|      |exact_match|↑  |0.7132|±  |0.0040|
| - biology         |    2.1|custom-extract|     5|exact_match|↑  |0.7127|±  |0.0169|
| - business        |    2.1|custom-extract|     5|exact_match|↑  |0.7643|±  |0.0151|
| - chemistry       |    2.1|custom-extract|     5|exact_match|↑  |0.7827|±  |0.0123|
| - computer_science|    2.1|custom-extract|     5|exact_match|↑  |0.6683|±  |0.0233|
| - economics       |    2.1|custom-extract|     5|exact_match|↑  |0.8211|±  |0.0132|
| - engineering     |    2.1|custom-extract|     5|exact_match|↑  |0.6182|±  |0.0156|
| - health          |    2.1|custom-extract|     5|exact_match|↑  |0.7323|±  |0.0155|
| - history         |    2.1|custom-extract|     5|exact_match|↑  |0.5958|±  |0.0252|
| - law             |    2.1|custom-extract|     5|exact_match|↑  |0.4705|±  |0.0150|
| - math            |    2.1|custom-extract|     5|exact_match|↑  |0.8179|±  |0.0105|
| - other           |    2.1|custom-extract|     5|exact_match|↑  |0.6872|±  |0.0153|
| - philosophy      |    2.1|custom-extract|     5|exact_match|↑  |0.5852|±  |0.0221|
| - physics         |    2.1|custom-extract|     5|exact_match|↑  |0.7852|±  |0.0114|
| - psychology      |    2.1|custom-extract|     5|exact_match|↑  |0.7757|±  |0.0148|