[Bugfix] fix kimi-linear crash

[ZJY0516]

Purpose

FIX #28436

TODO: it will crashes with tp=4, dp =2, IMA in fused_recurrent_gated_delta_rule_fwd_kernel(only in FULL_AND_PIECEWISE mode)

Test Plan

vllm serve moonshotai/Kimi-Linear-48B-A3B-Instruct --enable-expert-parallel -tp 8 --trust-remote-code
lm_eval --model local-completions --model_args "model=moonshotai/Kimi-Linear-48B-A3B-Instruct,base_url=http://0.0.0.0:8000/v1/completions,tokenized_requests=False,tokenizer_backend=None,num_concurrent=280" --tasks gsm8k --num_fewshot 5

Test Result

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.8901	±	0.0086
		strict-match	5	exact_match	↑	0.8779	±	0.0090

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[gemini-code-assist]

Code Review

This pull request appears to correctly fix a crash in the Kimi-Linear model by addressing issues with tensor shapes and indexing, particularly in speculative decoding scenarios. The changes involve removing an unused parameter g2, consistently using num_actual_tokens for slicing tensors to handle padding correctly, and making output assignments more robust. While these changes are positive, I've identified a potential out-of-bounds access bug in the decode path that could lead to a crash under certain conditions. Please see the detailed comment.

[jeejeelee]

The gsm8k result looks pretty bad, see: #27905

[ZJY0516]

The gsm8k result looks pretty bad, see: #27905

This is the result of main, which is almost same as this pr

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.8802	±	0.0089
		strict-match	5	exact_match	↑	0.3510	±	0.0131

[jeejeelee]

I will verify the main branch ASAP, also cc @zhiyuan1i

[jeejeelee]

I used the following script to test gsm8k

vllm serve moonshotai/Kimi-Linear-48B-A3B-Instruct --enable-expert-parallel -tp 8 --trust-remote-code

The metric result as follow:

local-completions (model=moonshotai/Kimi-Linear-48B-A3B-Instruct,base_url=http://0.0.0.0:8000/v1/completions,tokenized_requests=False,tokenizer_backend=None,num_concurrent=280), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: 1
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.8931|±  |0.0085|
|     |       |strict-match    |     5|exact_match|↑  |0.8726|±  |0.0092|

Maybe you can try lm_eval --model local-completions --model_args "model=moonshotai/Kimi-Linear-48B-A3B-Instruct,base_url=http://0.0.0.0:8000/v1/completions,tokenized_requests=False,tokenizer_backend=None,num_concurrent=280" --tasks gsm8k --num_fewshot 5

[ZJY0516]

Thanks. Now the result is normal

local-completions (model=moonshotai/Kimi-Linear-48B-A3B-Instruct,base_url=http://0.0.0.0:8000/v1/completions,tokenized_requests=False,tokenizer_backend=None,num_concurrent=280), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: 1
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.8886|±  |0.0087|
|     |       |strict-match    |     5|exact_match|↑  |0.8772|±  |0.0090|

[zhiyuan1i]

Thanks for pointing out, expert parallel hasn't been tested before

[heheda12345]

LGTM! For dp case, maybe you can double check the dummy run code path?

[ZJY0516]

maybe you can double check the dummy run code path?

It will not crash after merging from main.

[ZJY0516]

#28665 makes dp=2, tp=4 crash again :(

[heheda12345]

Can you help to check what has happened?

[heheda12345]

Also CC @tdoublep

[ZJY0516]

Can you help to check what has happened?

I've attempted to investigate this, but I'm not yet familiar with the internals of the vLLM scheduler.