Summary of Changes

Hello @farazkh80, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

I've implemented and integrated the TRTLLM Gen MLA Decode Kernel, a new attention backend specifically designed to optimize decode operations for NVIDIA's Blackwell architecture. This enhancement aims to improve performance on the latest GPU hardware. While optimized for decoding, prefill operations will seamlessly fall back to the FlashInfer MLA backend. I've also updated the relevant documentation and added comprehensive unit tests to ensure the stability and correctness of this new feature.

Highlights

• New TRTLLM MLA Attention Backend: I've integrated the TRTLLM Gen MLA Decode Kernel, which is specifically optimized for NVIDIA's Blackwell architecture (e.g., B200 GPUs). This will significantly enhance the performance of decode operations on these newer hardware platforms.
• Decode-Only Optimization with Prefill Fallback: The new TRTLLM MLA backend currently focuses solely on decode operations. For prefill operations, including those with multimodal inputs, the system will gracefully fall back to the existing FlashInfer MLA backend to ensure compatibility and functionality.
• Comprehensive Documentation Updates: I've updated the documentation to reflect the addition of the TRTLLM MLA backend, providing clear instructions on its usage and detailing its specific characteristics, such as its decode-only nature and Blackwell optimization.
• Configuration and Compatibility Constraints: The new backend enforces specific page_size requirements (32 or 64) and will issue a warning or adjust the size if an unsupported value is provided. Additionally, it does not currently support speculative decoding, and an error will be raised if attempted.
• New Backend Implementation File: I've added a new Python file, python/sglang/srt/layers/attention/trtllm_mla_backend.py, which encapsulates the core logic for the TRTLLM MLA backend, including memory management and interaction with FlashInfer's TRTLLM kernels.
• Dedicated Unit Tests: To ensure the robustness and correctness of the new backend, I've included a dedicated test file, python/sglang/test/attention/test_trtllm_mla_backend.py, with various unit tests covering basic functionality, output matching against reference implementations, and metadata handling.

Code Review

This pull request integrates the TRTLLM MLA decode kernel. The changes include the new backend implementation, integration into the model runner and server arguments, comprehensive unit tests, and documentation updates. My review focuses on the correctness and consistency of these changes. I've found some inconsistencies in the documentation regarding speculative decoding and multimodal support for the new backend, which contradict the implementation. I've also suggested an improvement to the new unit tests to make them more comprehensive. The core implementation of the new backend appears solid.