SIGN IN SIGN UP

FlashInfer: Kernel Library for LLM Serving

0 0 188 Python

feat: Add support for LoRa delta in MOE mxint4 x bf16, MXFP8 & BF16 to trtllm backend (#3153)

<!-- .github/pull_request_template.md -->

## ๐Ÿ“Œ Description

<!-- What does this PR do? Briefly describe the changes and why theyโ€™re
needed. -->

This PR exposes a LoRa delta parameter in the `trtllm_*_moe_routed` APIs
for MXINT4xBF16 and BF16. It also returns the result of FC1 (after
activation) for the consumption when calculating LoRa deltas for FC2.

The intended use-case is for frameworks to implement calculating
per-token LoRa deltas and have the trtllm MOE backed consume them.

Key points:

1. LoRa deltas are expected in expanded format: `[num_tokens, top_k,
intermediate_size * 2]`
2. LoRa deltas are applied before activation, and only support swiglu
activation
3. Using LoRa requires frameworks to calculate the routing. This
prevents issues where different top-k results are calculated
4. The framework should parse the expanded FC1 output according to the
returned permuted indices map


## ๐Ÿ” Related Issues

<!-- Link any related issues here -->

## ๐Ÿš€ 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

- [x] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [x] I have installed the hooks with `pre-commit install`.
- [x] 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](https://pre-commit.com/).

## ๐Ÿงช Tests

- [ ] Tests have been added or updated as needed.
- [ ] All tests are passing (`unittest`, etc.).

## Reviewer Notes

<!-- Optional: anything you'd like reviewers to focus on, concerns, etc.
-->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added LoRA GEMM1 delta support for routed Mixture of Experts
operations
  * Introduced MXINT4 block-scale MoE kernel implementation
* Enhanced MoE kernels with optional activation output return capability

* **Improvements**
* Enhanced bias-aware kernel configuration filtering for GEMM operations
  * Refined GEMM1 bias handling with row index permutation support
* Updated MoE kernel launchers with extended parameter plumbing for bias
modes

<!-- review_stack_entry_start -->

[![Review Change
Stack](https://storage.googleapis.com/coderabbit_public_assets/review-stack-in-coderabbit-ui.svg)](https://app.coderabbit.ai/change-stack/flashinfer-ai/flashinfer/pull/3153?utm_source=github_walkthrough&utm_medium=github&utm_campaign=change_stack)

<!-- review_stack_entry_end -->
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
D
Daniel Stokes committed
6f651b63fce45753f47c0ed7d7651a6daa2f1cf8
Parent: 18f4534
Committed by GitHub <noreply@github.com> on 5/21/2026, 8:19:16 PM