Motivation

Fix #15089
Make TRTLLM MHA workspace size configurable via SGLANG_FLASHINFER_WORKSPACE_SIZE environment variable instead of hardcoded 512 MB. This allows users to adjust workspace size for large batch sizes to avoid OOM errors.

Modifications

Added _get_trtllm_mha_workspace_size() helper function to handle workspace size configuration

Modified TRTLLMHAAttnBackend.init() to use configurable workspace size

Preserves user's environment variable setting even if it gets overridden by FlashInferAttnBackend.init()

Default value remains 512 MB for backward compatibility

Accuracy Tests

Benchmarking and Profiling

running
export SGLANG_FLASHINFER_WORKSPACE_SIZE=$((2048 * 1024 * 1024))
python -m sglang.launch_server
--model-path Qwen/Qwen2.5-7B-Instruct
--attention-backend trtllm_mha
--port 30000
(when adding log info)

unset SGLANG_FLASHINFER_WORKSPACE_SIZE
python -m sglang.launch_server
--model-path Qwen/Qwen2.5-7B-Instruct
--attention-backend trtllm_mha
--port 30000

export SGLANG_FLASHINFER_WORKSPACE_SIZE=$((1024 * 1024 * 1024))
python -m sglang.launch_server
--model-path Qwen/Qwen2.5-7B-Instruct
--attention-backend trtllm_mha
--port 30000

export SGLANG_FLASHINFER_WORKSPACE_SIZE=$((2048 * 1024 * 1024))
python -m sglang.launch_server
--model-path Qwen/Qwen2.5-7B-Instruct
--attention-backend trtllm_mha
--port 30000

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.
• Follow the SGLang code style guidance.
• Work with maintainers to merge your PR. See the PR Merge Process