[Issue]: TypeError: kv_last_page_lens needs to be <class 'torch.Tensor'> but got <class 'float'>

[ChangLiu0709]

Problem Description

I'm trying to use the updated MLA for MTP from the PR: #505

Background:

• Cloned recursively of the new repo.

• Re-built it by:
  python3 setup.py develop

• Launch the SGLang server using the command of sglang.launch_server and there is an error as below:

File "/sgl-workspace/sglang/aiter/aiter/jit/core.py", line 571, in check_args
raise TypeError(
TypeError: kv_last_page_lens needs to be <class 'torch.Tensor'> but got <class 'float'>

Wondering if anyone met this issue before and knows about the solution.

Operating System

Ubuntu 22.04.5 LTS

CPU

AMD EPYC 9654 96-Core Processor

GPU

AMD Instinct MI300X

ROCm Version

ROCm 6.3.0

ROCm Component

No response

Steps to Reproduce

No response

(Optional for Linux users) Output of /opt/rocm/bin/rocminfo --support

No response

Additional Information

No response

[valarLip]

the interface updated to support mtp, need some update on sglang side before call new aiter kernel

[ChangLiu0709]

the interface updated to support mtp, need some update on sglang side before call new aiter kernel

Oh I see that's why it does not fit for now. May I ask when the update will be ready?

[valarLip]

sgl-project/sglang#6838 looks like it's merged

[ChangLiu0709]

sgl-project/sglang#6838 looks like it's merged

Will have a test soon.

[ChangLiu0709]

sgl-project/sglang#6838 looks like it's merged

Hi @valarLip I have tried rebuilding the docker image based on the updated sglang repo, (merged sgl-project/sglang#6838), but the performance is still under expected. Wondering if you know any solutions or ideas to fix the bug of kv_last_page_lens?

[ChangLiu0709]

I've further checked the code and realized that this bug is triggered due to the newly introduced two parameters, qo_indptr and max_seqlen_q, in aiter/mla.py:

def mla_decode_fwd(
q,
kv_buffer,
o,
qo_indptr,
kv_indptr,
kv_indices,
kv_last_page_lens,
max_seqlen_q,
sm_scale=None, # 1.0 / (qk_head_dim**0.5)
logit_cap=0.0,
num_kv_splits=None, # for experts only!!!
): ...