[Bugfix] Fix illegal memory access

[elvischenv]

Motivation

Fixed illegal memory access issue in #12695 and flashinfer-ai/flashinfer#2034
Caused by #12524

Modifications

Accuracy Tests

Benchmarking and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[elvischenv]

Hang issue WAR: Increase max_token_num to allocate a larger workspace

[nvpohanh]

add an assert: assert input_tensor.shape[0] <= max_token_num ?

[nvpohanh]

should be covered by the code below:

    if input_tensor.shape[0] > max_token_num:
        logger.debug(
            "Input token(%d) is greater than max_token_num(%d), "
            "falling back to standard implementation",
            input_tensor.shape[0],
            max_token_num,
        )
        return None, None

[Fridge003]

@elvischenv Failed at GPTOSS Ci test, please have a look
https://github.com/sgl-project/sglang/actions/runs/19149049114/job/54736119046?pr=12758

[elvischenv]

Failed at GPTOSS Ci test, please have a look
https://github.com/sgl-project/sglang/actions/runs/19149049114/job/54736119046?pr=12758

This is DeepseekV2.

  File "/actions-runner/_work/sglang/sglang/python/sglang/srt/models/deepseek_v2.py", line 793, in forward_normal_dual_stream
    final_hidden_states += shared_output
RuntimeError: The size of tensor a (3584) must match the size of tensor b (7168) at non-singleton dimension 1

This failure seems also related to #12524. cc @merrymercy

[Fridge003]

@elvischenv Other PRs can pass dpsk test.
https://github.com/sgl-project/sglang/actions/runs/19153565252/job/54752242601

[Fridge003]

The gptoss CI test on B200 is passing here
https://github.com/sgl-project/sglang/actions/runs/19157818700/job/54765839477?pr=12758