flashinfer-python 0.6.12rc2__tar.gz → 0.6.13rc1__tar.gz

flashinfer_python-0.6.13rc1/PKG-INFO

@@ -0,0 +1,304 @@
+Metadata-Version: 2.4
+Name: flashinfer-python
+Version: 0.6.13rc1
+Summary: FlashInfer: Kernel Library for LLM Serving
+Author: FlashInfer team
+License-Expression: Apache-2.0
+Project-URL: Homepage, https://github.com/flashinfer-ai/flashinfer
+Requires-Python: <4.0,>=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: apache-tvm-ffi!=0.1.8,!=0.1.8.post0,<0.2,>=0.1.6
+Requires-Dist: click
+Requires-Dist: cuda-tile[tileiras]
+Requires-Dist: einops
+Requires-Dist: ninja
+Requires-Dist: numpy
+Requires-Dist: nvidia-cudnn-frontend>=1.13.0
+Requires-Dist: nvidia-cutlass-dsl>=4.5.0
+Requires-Dist: nvidia-ml-py
+Requires-Dist: packaging>=24.2
+Requires-Dist: requests
+Requires-Dist: tabulate
+Requires-Dist: torch
+Requires-Dist: tqdm
+Provides-Extra: cu12
+Requires-Dist: nvidia-cutlass-dsl>=4.5.0; extra == "cu12"
+Provides-Extra: cu13
+Requires-Dist: nvidia-cutlass-dsl[cu13]>=4.5.0; extra == "cu13"
+Provides-Extra: nvep
+Requires-Dist: cuda-python>=13.0; extra == "nvep"
+Dynamic: license-file
+
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" srcset="https://github.com/flashinfer-ai/web-data/blob/main/logo/FlashInfer-black-background.png?raw=true">
+    <img alt="FlashInfer" src="https://github.com/flashinfer-ai/web-data/blob/main/logo/FlashInfer-white-background.png?raw=true" width=55%>
+  </picture>
+</p>
+<h1 align="center">
+High-Performance GPU Kernels for Inference
+</h1>
+
+<p align="center">
+| <a href="https://docs.flashinfer.ai"><b>Documentation</b></a> | <a href="https://github.com/flashinfer-ai/flashinfer/releases/latest"><b>Latest Release</b></a> | <a href="https://flashinfer.ai"><b>Blog</b></a> | <a href="https://join.slack.com/t/flashinfer/shared_invite/zt-379wct3hc-D5jR~1ZKQcU00WHsXhgvtA"><b>Slack</b></a> |  <a href="https://github.com/orgs/flashinfer-ai/discussions"><b>Discussion Forum</b></a> |
+</p>
+
+[![Build Status](https://ci.tlcpack.ai/job/flashinfer-ci/job/main/badge/icon)](https://ci.tlcpack.ai/job/flashinfer-ci/job/main/)
+[![Documentation](https://github.com/flashinfer-ai/flashinfer/actions/workflows/build-doc.yml/badge.svg)](https://github.com/flashinfer-ai/flashinfer/actions/workflows/build-doc.yml)
+
+**FlashInfer** is a library and kernel generator for inference that delivers state-of-the-art performance across diverse GPU architectures. It provides unified APIs for attention, GEMM, and MoE operations with multiple backend implementations including FlashAttention-2/3, cuDNN, CUTLASS, and TensorRT-LLM.
+
+## Why FlashInfer?
+
+- **State-of-the-art Performance**: Optimized kernels for prefill, decode, and mixed batching scenarios
+- **Multiple Backends**: Automatically selects the best backend for your hardware and workload
+- **Modern Architecture Support**: Support for SM75 (Turing) and later (through Blackwell)
+- **Low-Precision Compute**: FP8 and FP4 quantization for attention, GEMM, and MoE operations
+- **Production-Ready**: CUDAGraph and torch.compile compatible for low-latency serving
+
+## Core Features
+
+### Attention Kernels
+- **Paged and Ragged KV-Cache**: Efficient memory management for dynamic batch serving
+- **Decode, Prefill, and Append**: Optimized kernels for all attention phases
+- **MLA Attention**: Native support for DeepSeek's Multi-Latent Attention
+- **Cascade Attention**: Memory-efficient hierarchical KV-Cache for shared prefixes
+- **Sparse Attention**: Block-sparse and variable block-sparse patterns
+- **POD-Attention**: Fused prefill+decode for mixed batching
+
+### GEMM & Linear Operations
+- **BF16 GEMM**: BF16 matrix multiplication for SM10.0+ GPUs.
+- **FP8 GEMM**: Per-tensor and groupwise scaling
+- **FP4 GEMM**: NVFP4 and MXFP4 matrix multiplication for Blackwell GPUs
+- **Grouped GEMM**: Efficient batched matrix operations for LoRA and multi-expert routing
+
+### Mixture of Experts (MoE)
+- **Fused MoE Kernels**
+- **Multiple Routing Methods**: DeepSeek-V3, Llama-4, and standard top-k routing
+- **Quantized MoE**: FP8 and FP4 expert weights with block-wise scaling
+
+### Sampling & Decoding
+- **Sorting-Free Sampling**: Efficient Top-K, Top-P, and Min-P without sorting
+- **Speculative Decoding**: Chain speculative sampling support
+
+### Communication
+- **AllReduce**: Custom implementations
+- **Multi-Node NVLink**: MNNVL support for multi-node inference
+- **NVSHMEM Integration**: For distributed memory operations
+
+### Other Operators
+- **RoPE**: LLaMA-style rotary position embeddings (including LLaMA 3.1)
+- **Normalization**: RMSNorm, LayerNorm, Gemma-style fused operations
+- **Activations**: SiLU, GELU with fused gating
+
+## GPU Support
+
+| Architecture | Compute Capability | Example GPUs |
+|--------------|-------------------|------|
+| Turing | SM 7.5 | T4, RTX 20 series |
+| Ampere | SM 8.0, 8.6 | A100, A10, RTX 30 series |
+| Ada Lovelace | SM 8.9 | L4, L40, RTX 40 series |
+| Hopper | SM 9.0 | H100, H200 |
+| Blackwell | SM 10.0, 10.3 | B200, B300 |
+| Blackwell | SM 11.0 | Jetson Thor |
+| Blackwell | SM 12.0, 12.1 | RTX 50 series, DGX Spark |
+
+> **Note:** Not all features are supported across all compute capabilities.
+
+## News
+
+Latest: [![GitHub Release](https://img.shields.io/github/v/release/flashinfer-ai/flashinfer)](https://github.com/flashinfer-ai/flashinfer/releases/latest)
+
+Notable updates:
+- [2025-10-08] Blackwell support added in [v0.4.0](https://github.com/flashinfer-ai/flashinfer/releases/tag/v0.4.0)
+- [2025-03-10] [Blog Post](https://flashinfer.ai/2025/03/10/sampling.html) Sorting-Free GPU Kernels for LLM Sampling, which explains the design of sampling kernels in FlashInfer.
+
+## Getting Started
+
+### Installation
+
+**Quickstart:**
+
+```bash
+pip install flashinfer-python
+```
+
+**Package Options:**
+
+- **flashinfer-python**: Core package that compiles/downloads kernels on first use
+- **flashinfer-cubin**: Pre-compiled kernel binaries for all supported GPU architectures
+- **flashinfer-jit-cache**: Pre-built kernel cache for specific CUDA versions
+
+**For faster initialization and offline usage**, install the optional packages to have most kernels pre-compiled:
+
+```bash
+pip install flashinfer-python flashinfer-cubin
+# JIT cache (replace cu129 with your CUDA version)
+pip install flashinfer-jit-cache --index-url https://flashinfer.ai/whl/cu129
+```
+
+**For Blackwell (SM100+) CuTe DSL kernels**, install with the CUDA 13 extra to enable Blackwell-optimized kernels:
+
+```bash
+pip install flashinfer-python[cu13]
+```
+
+### Verify Installation
+
+```bash
+flashinfer show-config
+```
+
+### Basic Usage
+
+```python
+import torch
+import flashinfer
+
+# Single decode attention
+q = torch.randn(32, 128, device="cuda", dtype=torch.float16)  # [num_qo_heads, head_dim]
+k = torch.randn(2048, 32, 128, device="cuda", dtype=torch.float16)  # [kv_len, num_kv_heads, head_dim]
+v = torch.randn(2048, 32, 128, device="cuda", dtype=torch.float16)
+
+output = flashinfer.single_decode_with_kv_cache(q, k, v)
+```
+
+See [documentation](https://docs.flashinfer.ai/) for comprehensive API reference and tutorials.
+
+### Install from Source
+
+```bash
+git clone https://github.com/flashinfer-ai/flashinfer.git --recursive
+cd flashinfer
+python -m pip install -v .
+```
+
+**For development**, install in editable mode:
+
+```bash
+python -m pip install --no-build-isolation -e . -v
+```
+
+> **Note:** When using `--no-build-isolation`, pip does not automatically install build dependencies. FlashInfer requires `setuptools>=77`. If you encounter an error like `AttributeError: module 'setuptools.build_meta' has no attribute 'prepare_metadata_for_build_editable'`, upgrade pip and setuptools first:
+> ```bash
+> python -m pip install --upgrade pip setuptools
+> ```
+
+Build optional packages:
+
+```bash
+# flashinfer-cubin
+cd flashinfer-cubin
+python -m build --no-isolation --wheel
+python -m pip install dist/*.whl
+```
+
+```bash
+# flashinfer-jit-cache (customize for your target GPUs)
+export FLASHINFER_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a 10.0a 10.3a 11.0a 12.0f"
+cd flashinfer-jit-cache
+python -m build --no-isolation --wheel
+python -m pip install dist/*.whl
+```
+
+For more details, see the [Install from Source documentation](https://docs.flashinfer.ai/installation.html#install-from-source).
+
+### Nightly Builds
+
+```bash
+pip install -U --pre flashinfer-python --index-url https://flashinfer.ai/whl/nightly/ --no-deps
+pip install flashinfer-python  # Install dependencies from PyPI
+pip install -U --pre flashinfer-cubin --index-url https://flashinfer.ai/whl/nightly/
+# JIT cache (replace cu129 with your CUDA version)
+pip install -U --pre flashinfer-jit-cache --index-url https://flashinfer.ai/whl/nightly/cu129
+```
+
+### CLI Tools
+
+FlashInfer provides several CLI commands for configuration, module management, and development:
+
+```bash
+# Verify installation and view configuration
+flashinfer show-config
+
+# List and inspect modules
+flashinfer list-modules
+flashinfer module-status
+
+# Manage artifacts and cache
+flashinfer download-cubin
+flashinfer clear-cache
+
+# For developers: generate compile_commands.json for IDE integration
+flashinfer export-compile-commands [output_path]
+```
+
+For complete documentation, see the [CLI reference](https://docs.flashinfer.ai/cli.html).
+
+## API Logging
+
+FlashInfer provides comprehensive API logging for debugging. Enable it using environment variables:
+
+```bash
+# Enable logging (levels: 0=off (default), 1=basic, 3=detailed, 5=statistics)
+export FLASHINFER_LOGLEVEL=3
+
+# Set log destination (stdout (default), stderr, or file path)
+export FLASHINFER_LOGDEST=stdout
+```
+
+For detailed information about logging levels, configuration, and advanced features, see [Logging](https://docs.flashinfer.ai/logging.html) in our documentation.
+
+## Custom Attention Variants
+
+Users can customize their own attention variants with additional parameters. For more details, refer to our [JIT examples](https://github.com/flashinfer-ai/flashinfer/blob/main/tests/utils/test_jit_example.py).
+
+## CUDA Support
+
+**Supported CUDA Versions:** 12.6, 12.8, 13.0, 13.1
+
+> **Note:** FlashInfer strives to follow PyTorch's supported CUDA versions plus the latest CUDA release.
+
+## Adoption
+
+FlashInfer powers inference in:
+
+- [SGLang](https://github.com/sgl-project/sglang)
+- [vLLM](https://github.com/vllm-project/vllm)
+- [TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM)
+- [TGI (Text Generation Inference)](https://github.com/huggingface/text-generation-inference)
+- [MLC-LLM](https://github.com/mlc-ai/mlc-llm)
+- [LightLLM](https://github.com/ModelTC/lightllm)
+- [lorax](https://github.com/predibase/lorax)
+- [ScaleLLM](https://github.com/vectorch-ai/ScaleLLM)
+
+## Acknowledgement
+
+FlashInfer is inspired by [FlashAttention](https://github.com/dao-AILab/flash-attention/), [vLLM](https://github.com/vllm-project/vllm), [stream-K](https://arxiv.org/abs/2301.03598), [CUTLASS](https://github.com/nvidia/cutlass), and [AITemplate](https://github.com/facebookincubator/AITemplate).
+
+## Citation
+
+If you find FlashInfer helpful in your project or research, please consider citing our [paper](https://arxiv.org/abs/2501.01005):
+
+```bibtex
+@article{ye2025flashinfer,
+    title = {FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving},
+    author = {
+      Ye, Zihao and
+      Chen, Lequn and
+      Lai, Ruihang and
+      Lin, Wuwei and
+      Zhang, Yineng and
+      Wang, Stephanie and
+      Chen, Tianqi and
+      Kasikci, Baris and
+      Grover, Vinod and
+      Krishnamurthy, Arvind and
+      Ceze, Luis
+    },
+    journal = {arXiv preprint arXiv:2501.01005},
+    year = {2025},
+    url = {https://arxiv.org/abs/2501.01005}
+}
+```

flashinfer_python-0.6.13rc1/csrc/bgmv_moe/kernel_config.h

@@ -0,0 +1,40 @@
+#pragma once
+
+/*
+ * BGMV MoE kernel tuning parameters.
+ *
+ * Target: H100/H200 (sm_90, 228 KB shared memory per SM)
+ * Also supports sm_70+ (V100, A100) with reduced pipeline depth.
+ *
+ * Copyright (c) 2025 by FlashInfer team.
+ * Licensed under the Apache License, Version 2.0.
+ */
+
+struct MoeShrinkKernelConfig {
+  static constexpr int tx = 32;        // threads per warp (x-dimension)
+  static constexpr int ty = 4;         // number of warps (y-dimension)
+  static constexpr int vec_size = 8;   // elements per vectorized load
+  static constexpr int rank_tile = 8;  // rank elements per block (8x X reuse)
+
+  // Multi-pair decode path: PPB=4 pairs per block for decode,
+  // PPB=1 for prefill (grid already saturates GPU).
+  static constexpr int pairs_per_block_prefill = 1;
+  static constexpr int pairs_per_block_decode = 4;
+  static constexpr int decode_threshold = 32;
+
+  // Pipeline depth: 3 stages on sm_90 decode (216 KB / 228 KB = 95%),
+  // 2 stages for prefill (36 KB, leaves room for occupancy).
+  static constexpr int num_stages_default = 2;
+  static constexpr int num_stages_extended = 3;
+
+  // Shared memory budget (decode, PPB=4, 3 stages, RANK_TILE=8, fp16):
+  //   X: 3 * 4 * 1024 * 2 =  24 KB
+  //   W: 3 * 4 * 8 * 1024 * 2 = 192 KB
+  //   y: 4 * 8 * 4 * 4 =  512 B
+  //   Total: ~216 KB (fits 228 KB on H100/H200)
+};
+
+struct MoeExpandKernelConfig {
+  static constexpr int tz = 4;
+  static constexpr int vec_size = 8;
+};

flashinfer_python-0.6.13rc1/csrc/bgmv_moe/moe_bgmv_bf16_bf16_bf16.cu

@@ -0,0 +1,10 @@
+/*
+ * Copyright (c) 2025 by FlashInfer team.
+ * Licensed under the Apache License, Version 2.0.
+ */
+
+#include "moe_bgmv_config.h"
+#include "moe_bgmv_impl.cuh"
+
+// Shrink + expand (in_T=out_T=W_T=nv_bfloat16).
+FOR_MOE_ALL_WIDE_NARROW(INST_MOE_BGMV_TWOSIDE, nv_bfloat16, nv_bfloat16, nv_bfloat16)

flashinfer_python-0.6.13rc1/csrc/bgmv_moe/moe_bgmv_bf16_fp32_bf16.cu

@@ -0,0 +1,15 @@
+/*
+ * Copyright (c) 2025 by FlashInfer team.
+ * Licensed under the Apache License, Version 2.0.
+ */
+
+#include "moe_bgmv_config.h"
+#include "moe_bgmv_impl.cuh"
+
+// Expand only (in_T=nv_bfloat16, Y=float32, W_T=nv_bfloat16).
+// Shrink is covered by moe_bgmv_bf16_bf16_bf16.cu.
+
+#define INST_MOE_BGMV_EXPAND_ONLY(in_T, out_T, W_T, narrow, wide) \
+  INST_MOE_BGMV_EXPAND_SLICED(narrow, wide, in_T, W_T)
+
+FOR_MOE_ALL_WIDE_NARROW(INST_MOE_BGMV_EXPAND_ONLY, nv_bfloat16, float, nv_bfloat16)

flashinfer_python-0.6.13rc1/csrc/bgmv_moe/moe_bgmv_binding.cu

@@ -0,0 +1,23 @@
+/*
+ * TVM-FFI binding for BGMV MoE kernels.
+ *
+ * Exports two functions:
+ *   - bgmv_moe_shrink
+ *   - bgmv_moe_expand
+ *
+ * Copyright (c) 2025 by FlashInfer team.
+ * Licensed under the Apache License, Version 2.0.
+ */
+
+#include "moe_bgmv_ops.cu"
+
+// Forward declarations
+void bgmv_moe_shrink(TensorView y, TensorView x, TensorView w_ptr, TensorView sorted_token_ids,
+                     TensorView expert_ids, TensorView lora_indices, int64_t lora_stride);
+
+void bgmv_moe_expand(TensorView y, TensorView x, TensorView w_ptr, TensorView sorted_token_ids,
+                     TensorView expert_ids, TensorView topk_weights, TensorView lora_indices,
+                     TensorView slice_start_loc, int64_t first_feat_out, int64_t lora_stride);
+
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(bgmv_moe_shrink, bgmv_moe_shrink);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(bgmv_moe_expand, bgmv_moe_expand);

flashinfer_python-0.6.13rc1/csrc/bgmv_moe/moe_bgmv_config.h

@@ -0,0 +1,82 @@
+#pragma once
+
+#include <cstdint>
+
+/*
+ * BGMV MoE dimension configuration and forward declarations.
+ *
+ * Defines the set of (narrow, wide) dimension pairs that are compiled.
+ * narrow = LoRA rank (8, 16, 32, 64)
+ * wide = hidden/intermediate dimension of the model
+ *
+ * Models covered:
+ *   Qwen3-30B-A3B:              gate_up=(2048,768),  down=(768,2048)
+ *   Qwen3.5-35B-A3B:            gate_up=(2048,768),  down=(768,2048)  [256 experts, top-8]
+ *   Gemma-4-26B-A4B:            gate_up=(2816,2112), down=(2112,2816) [128 experts, top-8]
+ *   Nemotron-Nano-3-30B-A3B:    gate_up=(2688,1856), down=(1856,2688)
+ *   Nemotron-3-Super-120B-A12B: gate_up=(4096,2688), down=(2688,4096)
+ *   Large MoE (128 experts):    gate_up=(3072,5888), down=(2944,3072)
+ *
+ * TP-sharded dimensions also included.
+ * All wide values must satisfy: wide % 32 == 0.
+ *
+ * Copyright (c) 2025 by FlashInfer team.
+ * Licensed under the Apache License, Version 2.0.
+ */
+
+// clang-format off
+
+#define FOR_MOE_ALL_WIDE(f, in_T, out_T, W_T, narrow) \
+    f(in_T, out_T, W_T, narrow, 384)   \
+    f(in_T, out_T, W_T, narrow, 736)   \
+    f(in_T, out_T, W_T, narrow, 768)   \
+    f(in_T, out_T, W_T, narrow, 1024)  \
+    f(in_T, out_T, W_T, narrow, 1344)  \
+    f(in_T, out_T, W_T, narrow, 1472)  \
+    f(in_T, out_T, W_T, narrow, 1536)  \
+    f(in_T, out_T, W_T, narrow, 1856)  \
+    f(in_T, out_T, W_T, narrow, 2048)  \
+    f(in_T, out_T, W_T, narrow, 2112)  \
+    f(in_T, out_T, W_T, narrow, 2688)  \
+    f(in_T, out_T, W_T, narrow, 2816)  \
+    f(in_T, out_T, W_T, narrow, 2880)  \
+    f(in_T, out_T, W_T, narrow, 2944)  \
+    f(in_T, out_T, W_T, narrow, 3072)  \
+    f(in_T, out_T, W_T, narrow, 4096)  \
+    f(in_T, out_T, W_T, narrow, 5120)  \
+    f(in_T, out_T, W_T, narrow, 5888)  \
+    f(in_T, out_T, W_T, narrow, 7168)  \
+    f(in_T, out_T, W_T, narrow, 8192)  \
+    f(in_T, out_T, W_T, narrow, 10240) \
+    f(in_T, out_T, W_T, narrow, 14336) \
+    f(in_T, out_T, W_T, narrow, 16384) \
+    f(in_T, out_T, W_T, narrow, 28672)
+
+#define FOR_MOE_ALL_WIDE_NARROW(f, in_T, out_T, W_T) \
+    FOR_MOE_ALL_WIDE(f, in_T, out_T, W_T, 8)  \
+    FOR_MOE_ALL_WIDE(f, in_T, out_T, W_T, 16) \
+    FOR_MOE_ALL_WIDE(f, in_T, out_T, W_T, 32) \
+    FOR_MOE_ALL_WIDE(f, in_T, out_T, W_T, 64)
+
+// clang-format on
+
+// ===== Forward declarations =====
+
+template <int feat_in, int feat_out, typename in_T, typename out_T, typename W_T>
+void moe_bgmv_shrink_sliced(out_T* __restrict__ Y, const in_T* __restrict__ X,
+                            W_T** __restrict__ w_ptr, const int64_t* __restrict__ sorted_token_ids,
+                            const int64_t* __restrict__ expert_ids,
+                            const int64_t* __restrict__ lora_indices, int64_t num_pairs,
+                            int64_t num_slices, int64_t num_experts, int64_t num_tokens,
+                            int64_t lora_stride, float scale);
+
+template <int feat_in, int feat_out, typename in_T, typename W_T>
+void moe_bgmv_expand_sliced(float* __restrict__ Y, const in_T* __restrict__ X,
+                            W_T** __restrict__ w_ptr, const int64_t* __restrict__ sorted_token_ids,
+                            const int64_t* __restrict__ expert_ids,
+                            const int64_t* __restrict__ lora_indices,
+                            const float* __restrict__ topk_weights,
+                            const int64_t* __restrict__ slice_start_loc, int64_t num_pairs,
+                            int64_t num_slices, int64_t num_experts, int64_t total_feat_out,
+                            int32_t current_feat_out, int64_t num_tokens, int64_t lora_stride,
+                            float scale);

flashinfer_python-0.6.13rc1/csrc/bgmv_moe/moe_bgmv_fp16_fp16_fp16.cu

@@ -0,0 +1,10 @@
+/*
+ * Copyright (c) 2025 by FlashInfer team.
+ * Licensed under the Apache License, Version 2.0.
+ */
+
+#include "moe_bgmv_config.h"
+#include "moe_bgmv_impl.cuh"
+
+// Shrink + expand (in_T=out_T=W_T=nv_half).
+FOR_MOE_ALL_WIDE_NARROW(INST_MOE_BGMV_TWOSIDE, nv_half, nv_half, nv_half)

flashinfer_python-0.6.13rc1/csrc/bgmv_moe/moe_bgmv_fp16_fp32_fp16.cu

@@ -0,0 +1,15 @@
+/*
+ * Copyright (c) 2025 by FlashInfer team.
+ * Licensed under the Apache License, Version 2.0.
+ */
+
+#include "moe_bgmv_config.h"
+#include "moe_bgmv_impl.cuh"
+
+// Expand only (in_T=nv_half, Y=float32, W_T=nv_half).
+// Shrink is covered by moe_bgmv_fp16_fp16_fp16.cu.
+
+#define INST_MOE_BGMV_EXPAND_ONLY(in_T, out_T, W_T, narrow, wide) \
+  INST_MOE_BGMV_EXPAND_SLICED(narrow, wide, in_T, W_T)
+
+FOR_MOE_ALL_WIDE_NARROW(INST_MOE_BGMV_EXPAND_ONLY, nv_half, float, nv_half)

flashinfer_python-0.6.13rc1/csrc/bgmv_moe/moe_bgmv_fp32_bf16_bf16.cu

@@ -0,0 +1,15 @@
+/*
+ * Copyright (c) 2025 by FlashInfer team.
+ * Licensed under the Apache License, Version 2.0.
+ */
+
+#include "moe_bgmv_config.h"
+#include "moe_bgmv_impl.cuh"
+
+// Shrink only with mixed precision (in_T=float, out_T=nv_bfloat16, W_T=nv_bfloat16).
+// This handles the case where X is accumulated in fp32 but LoRA weights are bf16.
+
+#define INST_MOE_BGMV_SHRINK_ONLY(in_T, out_T, W_T, narrow, wide) \
+  INST_MOE_BGMV_SHRINK_SLICED(wide, narrow, in_T, out_T, W_T)
+
+FOR_MOE_ALL_WIDE_NARROW(INST_MOE_BGMV_SHRINK_ONLY, float, nv_bfloat16, nv_bfloat16)

flashinfer_python-0.6.13rc1/csrc/bgmv_moe/moe_bgmv_fp32_fp16_fp16.cu

@@ -0,0 +1,14 @@
+/*
+ * Copyright (c) 2025 by FlashInfer team.
+ * Licensed under the Apache License, Version 2.0.
+ */
+
+#include "moe_bgmv_config.h"
+#include "moe_bgmv_impl.cuh"
+
+// Shrink only with mixed precision (in_T=float, out_T=nv_half, W_T=nv_half).
+
+#define INST_MOE_BGMV_SHRINK_ONLY(in_T, out_T, W_T, narrow, wide) \
+  INST_MOE_BGMV_SHRINK_SLICED(wide, narrow, in_T, out_T, W_T)
+
+FOR_MOE_ALL_WIDE_NARROW(INST_MOE_BGMV_SHRINK_ONLY, float, nv_half, nv_half)