fix: semaphoress must be at the fixed range in workspace buffer on trtllm_gen attention

[yyihuang]

📌 Description

workspace_buffer arrangement

• on main branch
  paged_attention: counter (fixed 8MB at head) | scratch
  ragged_attention: softmax | counter (8MB, but not at the fixed address range) | scratch

• on PR branch:
  counter (fixed 8MB at head) | softmax | scratch

The range of semaphores must be fixed across multiple execution, since we are not clearing the buffer by zeros explicitly any more.

related PR: #1463

And #1566 (WIP) depends on this.

🔍 Related Issues

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

• Tests have been added or updated as needed.
• All tests are passing (unittest, etc.).

Reviewer Notes

[gemini-code-assist]

Summary of Changes

Hello @yyihuang, 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!

This pull request implements a crucial fix for the trtllm_gen attention mechanism by optimizing the memory layout of semaphores within the workspace buffer. The change involves moving these semaphores from the initial part of the buffer to its very end, ensuring they occupy a dedicated and non-conflicting space. This adjustment, coupled with the explicit definition of a 128MB maximum workspace size, enhances memory management and stability for attention operations in TensorRT-LLM.

Highlights

• Memory Layout Optimization: This pull request addresses a memory layout issue by relocating the semaphores used in trtllm_gen attention mechanisms. Previously, these semaphores were placed at the beginning of the workspace buffer, which could lead to conflicts or inefficient memory usage. They are now moved to the end of a fixed 128MB workspace buffer.
• Workspace Buffer Definition: A new constexpr constant, kMaxWorkspaceBufferSize, has been introduced and set to 128MB. This explicitly defines the maximum size of the workspace buffer for trtllm_gen operations, providing clarity and consistency in memory allocation.
• Pointer Calculation Adjustment: The pointer calculations for multiCtasKvScratchPtr within both trtllm_paged_attention_launcher and trtllm_ragged_attention_launcher functions have been updated to reflect the new semaphore placement at the tail of the workspace buffer.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request aims to refactor the workspace buffer management for the TRT-LLM generation attention mechanism by relocating the semaphore buffer to the tail of the workspace. This is accomplished by introducing a kMaxWorkspaceBufferSize constant and updating the pointer arithmetic accordingly. The change improves the clarity of the memory layout and appears to resolve a potential memory overlap issue in the trtllm_ragged_attention_launcher. My review includes one critical comment: essential C++ standard library headers have been commented out, which will lead to a compilation failure. This must be fixed for the changes to be functional.

[gemini-code-assist]

Commenting out these headers will cause a compilation error. The TllmGenFmhaRunnerCache class uses std::shared_ptr (from <memory>) and std::mutex with std::lock_guard (from <mutex>). These headers are necessary for the code to compile and should be restored.

#include <memory>
#include <mutex>

[yzh119]

Can you share the motivation of this change?

Also cc @elfiegg for visibility

[yzh119]

I would encourage to pass in this value as function arguments at python instead of hardcoded, to get rid of the possible issues of any changes on python side in the future.

[yyihuang]

Sure. We could do it in our future PR.