-
-
Notifications
You must be signed in to change notification settings - Fork 22.7k
[Kernel] Add native B12X linear, MoE, and causal attention backends #51696
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Closed
lukealonso
wants to merge
5
commits into
vllm-project:main
from
local-inference-lab:dev/b12x-moe-attn-lin
Closed
Changes from all commits
Commits
Show all changes
5 commits
Select commit
Hold shift + click to select a range
1d8a5c2
feat(linear): add B12X FP8 backends
lukealonso a5b529a
feat(linear): add B12X FP4 backends
lukealonso 9b3afae
feat(moe): add B12X FP4 backend
lukealonso 1697565
feat(attention): add B12X causal paged backend
lukealonso 9faf6bd
docs: document B12X backend usage
lukealonso File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,52 @@ | ||
| # B12X Backends | ||
|
|
||
| [B12X](https://pypi.org/project/b12x/) provides optional CUDA kernels for | ||
| NVIDIA SM120 and SM121 GPUs. Install the dependency with: | ||
|
|
||
| ```bash | ||
| uv pip install "vllm[b12x]" | ||
| ``` | ||
|
|
||
| The B12X backends are opt-in. To select all supported B12X kernels explicitly: | ||
|
|
||
| ```bash | ||
| vllm serve <model> \ | ||
| --linear-backend b12x \ | ||
| --moe-backend b12x \ | ||
| --attention-backend B12X_ATTN | ||
| ``` | ||
|
|
||
| Only pass `--moe-backend b12x` for a compatible NVFP4 or MXFP4 MoE model. The | ||
| linear and attention backends can be selected independently. | ||
|
|
||
| B12X uses MXFP8 activations by default for MXFP4 MoE and the checkpoint's | ||
| activation format for NVFP4 MoE. MXFP4 falls back to BF16 when its A8 path does | ||
| not support the model configuration. Set `VLLM_B12X_MOE_FORCE_A16=1` to force | ||
| BF16 activations for either weight format. | ||
|
|
||
| For more targeted linear selection, leave `--linear-backend` at `auto` and | ||
| enable B12X for one or both quantization families: | ||
|
|
||
| ```bash | ||
| VLLM_USE_B12X_FP8_GEMM=1 VLLM_USE_B12X_FP4_GEMM=1 \ | ||
| vllm serve <model> | ||
| ``` | ||
|
|
||
| An explicit non-B12X `--linear-backend` takes precedence over these environment | ||
| variables. | ||
|
|
||
| ## Supported Configurations | ||
|
|
||
| | Backend | Supported configurations | | ||
| | ------- | ------------------------ | | ||
| | Linear | Per-tensor FP8, 128x128 block FP8, MXFP8, NVFP4, and MXFP4 | | ||
| | MoE | Tensor-parallel MXFP4 weights with BF16 or MXFP8 activations; NVFP4 weights with BF16, NVFP4, or MXFP8 activations | | ||
| | Attention | Causal decoder MHA/MQA/GQA with BF16 model dtype; BF16, FP16, or FP8 E4M3 KV cache; head sizes 64, 128, 192, or 256; block size 64 or 128 | | ||
|
|
||
| The attention backend supports decode, prefill, mixed batches, speculative | ||
| verification, sliding-window attention, attention sinks, and CUDA graphs. | ||
|
|
||
| The B12X MoE backend does not support expert parallelism, expert maps, EXL3, or | ||
| NF3. The attention backend does not support MLA, non-causal attention, ALiBi, | ||
| logits soft capping, or context parallelism. Dense W4A16 layers are not handled | ||
| by B12X and continue to use another compatible backend such as Marlin. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think we should pin the version and update as we need to, like flashinfer