[Feature]: DeepSeek V3.2 prefill is 2x slower than V3.1

[Shang-Pin]

🚀 The feature, motivation and pitch

We are running DeepSeek V3.2 in TensorRT-LLM. We benchmarked the speed of prefills for both v3 and v3.2:

def test():
...   t = time.time()
...   text = "Hello " * 150000
...   openai.completions.create(model="", prompt="\n\n<|User|>%s    <|Assistant|>      **</think>**" % text, max_tokens=100)
...   print(time.time() - t)

DeepSeek V3.1 takes 8 seconds
DeepSeek V3.2 takes 18 seconds

Alternatives

The sparse attention is supposed to be faster or at least the same for prefill.

Additional context

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.

[chang-l]

Hi @Shang-Pin ,I see you set max_tokens=100, so enabled chunked prefill is enabled. Could you share the platform (GPU arch) you tested on?
We’re working on optimizations for the DSV3.2 indexer model for this scenario—stay tuned. In the meantime, can you please try disabling chunked prefill (increasing max_tokens>seq_len) for long sequences and see if that improves performance?

[Shang-Pin]

@chang-l We are testing on B200.
We are testing with max_num_tokens=8196 because of this issue.
The max_tokens=100 in the request is we only ask for maximum of 100 output tokens. That's just to make sure the majority of time is spent on prefill.
I can test without chunked prefill to verify if it's faster, however, we really need to enable chunked prefill for our production traffic. (max_seq_len can be up to 163840)

[chang-l]

@Shang-Pin Understood.

For this issue, I have a fix PR here: #9968

Also, considering your use case (max_seq_len can be up to 163840), you don't necessarily need chunked prefill. You can:

• Set max_seq_len/max_num_tokens = 163840(chunked prefill default off), and add the following section to your yml file:

sparse_attention_config:
    algorithm: dsa
    indexer_max_chunk_size: 8196 # (default is 32768; reducing this lowers the memory footprint quadratically)

This should alleviate the memory pressure during indexer logit computation. My local test shows it can support max_seq_len / max_num_tokens = 163840.

• We’ll continue improving chunked prefill performance.

Edit: btw, you will need the fix: #9968 for both cases.

[chang-l]

Looping in @lfr-0531