Port TRT-LLM communication kernels to flashinfer

[djns99]

📌 Description

This ports the latest MNNVL A2A communication implementation from TRT-LLM

🔍 Related Issues

#2094

🚀 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

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

✨ Finishing touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[trevor-m]

#1245 adds an optional custom communicator for the MnnvlMemory initialization. This is necessary because sglang/vllm don't use MPI so the previous bootstrapping wouldn't work. We will need to allow users to pass through the custom communicator here as well.

1. Can you add an arg to this function

config: Optional[MnnvlConfig] = None

2. Then after MnnvlMemory.initialize(), add

if config:
            MnnvlMemory.set_comm_from_config(mapping, config)  # type: ignore[attr-defined]

I think that should take care of it. Thanks!