Repository: huggingface/transformers
Pull request: [Performance] FP8 Grouped and Batched Matmuls

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

up to 30x faster than current fp8 experts, the kernels are also tailored for full torch.compile and cuda graphs compatibility.

```
============================================================
FP8Expert parity: eager / batched_mm / grouped_mm
============================================================

[case 1/5]
device=cuda  batch_size=1  num_tokens=64  total_tokens=64  num_experts=8  hidden=256  intermediate=512  top_k=2
  [eager vs eager_fused]  max=0.000000  mean=0.000000  PASS ✓
  [eager vs grouped_mm]  max=0.000000  mean=0.000000  PASS ✓
  [eager vs batched_mm]  max=0.000000  mean=0.000000  PASS ✓
  [batched_mm vs grouped_mm]  max=0.000000  mean=0.000000  PASS ✓

[case 2/5]
device=cuda  batch_size=1  num_tokens=1  total_tokens=1  num_experts=8  hidden=256  intermediate=512  top_k=2
  [eager vs eager_fused]  max=0.000000  mean=0.000000  PASS ✓
  [eager vs grouped_mm]  max=0.000000  mean=0.000000  PASS ✓
  [eager vs batched_mm]  max=0.000000  mean=0.000000  PASS ✓
  [batched_mm vs grouped_mm]  max=0.000000  mean=0.000000  PASS ✓

[case 3/5]
device=cuda  batch_size=1  num_tokens=7  total_tokens=7  num_experts=8  hidden=256  intermediate=512  top_k=1
  [eager vs eager_fused]  max=0.000000  mean=0.000000  PASS ✓
  [eager vs grouped_mm]  max=0.000000  mean=0.000000  PASS ✓
  [eager vs batched_mm]  max=0.000000  mean=0.000000  PASS ✓
  [batched_mm vs grouped_mm]  max=0.000000  mean=0.000000  PASS ✓

[case 4/5]
device=cuda  batch_size=1  num_tokens=4  total_tokens=4  num_experts=8  hidden=256  intermediate=512  top_k=8
  [eager vs eager_fused]  max=0.000000  mean=0.000000  PASS ✓
  [eager vs grouped_mm]  max=0.000000  mean=0.000000  PASS ✓
  [eager vs batched_mm]  max=0.000000  mean=0.000000  PASS ✓
  [batched_mm vs grouped_mm]  max=0.000000  mean=0.000000  PASS ✓

[case 5/5]
device=cuda  batch_size=4  num_tokens=64  total_tokens=256  num_experts=8  hidden=256  intermediate=512  top_k=2
  [eager vs eager_fused]  max=0.000000  mean=0.000000  PASS ✓
  [eager vs grouped_mm]  max=0.000000  mean=0.000000  PASS ✓
  [eager vs batched_mm]  max=0.000000  mean=0.000000  PASS ✓
  [batched_mm vs grouped_mm]  max=0.000000  mean=0.000000  PASS ✓

============================================================
All parity checks PASSED ✓

============================================================
Benchmark sweep
============================================================

────────────────────────────────────────────────────────────
Benchmark  device=cuda  batch_size=1  tokens=1  total=1  experts=8
           hidden=256  intermediate=512  top_k=2
────────────────────────────────────────────────────────────
  impl                            median (ms)    p10 (ms)    p90 (ms)   speedup
  eager                                 1.477       1.458       1.503  (baseline)
  eager_fused                           1.206       1.193       1.238     1.22x
  grouped_mm                            0.830       0.808       0.858     1.78x
  batched_mm                            0.423       0.418       0.441     3.49x
  eager (compiled)                      1.413       1.396       1.465     1.05x
  grouped_mm (compiled)                 0.138       0.136       0.147    10.67x
  batched_mm (compiled)                 0.138       0.136       0.151    10.71x

────────────────────────────────────────────────────────────
Benchmark  device=cuda  batch_size=1  tokens=8  total=8  experts=8
           hidden=256  intermediate=512  top_k=2
────────────────────────────────────────────────────────────
  impl                            median (ms)    p10 (ms)    p90 (ms)   speedup
  eager                                 4.506       4.430       4.581  (baseline)
  eager_fused                           3.672       3.654       3.700     1.23x
  grouped_mm                            0.809       0.797       0.840     5.57x
  batched_mm                            0.423       0.419       0.444    10.64x
  eager (compiled)                      4.752       4.733       4.775     0.95x
  grouped_mm (compiled)                 0.158       0.155       0.168    28.52x
  batched_mm (compiled)                 0.153       0.151       0.163    29.41x

────────────────────────────────────────────────────────────
Benchmark  device=cuda  batch_size=1  tokens=32  total=32  experts=8
           hidden=256  intermediate=512  top_k=2
────────────────────────────────────────────────────────────
  impl                            median (ms)    p10 (ms)    p90 (ms)   speedup
  eager                                 5.106       5.064       5.142  (baseline)
  eager_fused                           4.065       4.052       4.086     1.26x
  grouped_mm                            0.815       0.807       0.844     6.27x
  batched_mm                            0.433       0.428       0.455    11.80x
  eager (compiled)                      5.346       5.326       5.379     0.95x
  grouped_mm (compiled)                 0.159       0.157       0.167    32.01x
  batched_mm (compiled)                 0.167       0.165       0.177    30.62x

────────────────────────────────────────────────────────────
Benchmark  device=cuda  batch_size=1  tokens=128  total=128  experts=8
           hidden=256  intermediate=512  top_k=2
────────────────────────────────────────────────────────────
  impl                            median (ms)    p10 (ms)    p90 (ms)   speedup
  eager                                 5.236       5.098       5.275  (baseline)
  eager_fused                           4.094       4.073       4.197     1.28x
  grouped_mm                            0.826       0.815       0.854     6.34x
  batched_mm                            0.441       0.436       0.460    11.87x
  eager (compiled)                      5.572       5.545       5.611     0.94x
  grouped_mm (compiled)                 0.183       0.180       0.193    28.65x
  batched_mm (compiled)                 0.284       0.282       0.291    18.42x

────────────────────────────────────────────────────────────
Benchmark  device=cuda  batch_size=1  tokens=512  total=512  experts=8
           hidden=256  intermediate=512  top_k=2
────────────────────────────────────────────────────────────
  impl                            median (ms)    p10 (ms)    p90 (ms)   speedup
  eager                                 5.341       5.305       5.375  (baseline)
  eager_fused                           4.215       4.199       4.233     1.27x
  grouped_mm                            0.822       0.813       0.848     6.50x
  batched_mm                            1.405       1.391       1.416     3.80x
  eager (compiled)                      5.622       5.563       5.674     0.95x
  grouped_mm (compiled)                 0.219       0.217       0.224    24.40x
  batched_mm (compiled)                 0.724       0.712       0.734     7.38x

============================================================
```
<img width="4200" height="3000" alt="generate" src="https://github.com/user-attachments/assets/e47b74fc-d63f-4578-92d1-54d8095ac6a5" />

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker @Cyrilvallez
- vision models: @yonigozlan @molbap
- audio models: @eustlb @ebezzam @vasqu
- multimodal models: @zucchini-nlp
- graph models: @clefourrier

Library:

- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- continuous batching: @remi-or @ArthurZucker @McPatate
- pipelines: @Rocketknight1
- tokenizers: @ArthurZucker and @itazap
- trainer: @SunMarc
- attention: @vasqu @ArthurZucker @CyrilVallez
- model loading (from pretrained, etc): @CyrilVallez
- distributed: @3outeille @ArthurZucker
- CIs: @ydshieh

Integrations:

- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization: @SunMarc @MekkCyber
- kernels: @MekkCyber @drbh
- peft: @BenjaminBossan @githubnemo

Devices/Backends:

- AMD ROCm: @ivarflakstad
- Intel XPU: @IlyasMoutawwakil
- Ascend NPU: @ivarflakstad 

Documentation: @stevhliu

Research projects are not maintained and should be taken as is.

 -->


Changed files:
filename: src/transformers/integrations/finegrained_fp8.py
patch: @@ -11,51 +11,85 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import torch
+import torch.nn as nn
+import triton
+from torch.nn import functional as F
 
+from ..activations import ACT2FN
 from ..core_model_loading import ConversionOps
 from ..quantizers.quantizers_utils import should_convert_module
-from ..utils import is_kernels_available, is_torch_accelerator_available, is_torch_available, logging
+from ..utils import is_kernels_available, is_torch_available, logging
+from .hub_kernels import get_kernel
+from .moe import ExpertsInterface, use_experts_implementation
 
 
-if is_torch_available():
-    import torch
-    import torch.nn as nn
-    import triton
-    import triton.language as tl
-    from torch.nn import functional as F
+logger = logging.get_logger(__name__)
 
 
-logger = logging.get_logger(__name__)
+_FP8_DTYPE = torch.float8_e4m3fn
+_FP8_MIN = torch.finfo(_FP8_DTYPE).min
+_FP8_MAX = torch.finfo(_FP8_DTYPE).max
+
+# Global for the Triton quantization kernel (lazily compiled)
+_triton_kernel = None
+_triton_kernel_available = None
+
 
 # Global for the CUTLASS quantization kernel (lazily loaded)
-_quantization_kernel = None
+_cutlass_kernel = None
+_cutlass_kernel_available = None
+
 
+def _get_triton_kernel():
+    """Lazily compile the Triton quantization kernel."""
+    global _triton_kernel, _triton_kernel_available
 
-def _get_quantization_kernel():
+    if _triton_kernel_available is None:
+        # this kernel is universal so should be usable independently of torch version
+        _triton_kernel = get_kernel("kernels-community/finegrained-fp8")
+        _triton_kernel_available = True
+
+    return _triton_kernel
+
+
+def _get_cutlass_kernel():
     """Lazily load the CUTLASS quantization kernel from HuggingFace Hub."""
-    global _quantization_kernel
-    if _quantization_kernel is None:
-        try:
-            from .hub_kernels import get_kernel
+    global _cutlass_kernel, _cutlass_kernel_available
 
-            _quantization_kernel = get_kernel("RedHatAI/quantization")
+    if _cutlass_kernel_available is None:
+        try:
+            # this kernel's build was not updated since torch 2.8
+            _cutlass_kernel = get_kernel("RedHatAI/quantization")
+            _cutlass_kernel_available = True
         except Exception as e:
             logger.warning_once(f"Failed to load CUTLASS quantization kernel: {e}. Falling back to Triton.")
-            _quantization_kernel = False  # Mark as unavailable
-    return _quantization_kernel if _quantization_kernel else None
+            _cutlass_kernel_available = False
+
+    return _cutlass_kernel
 
 
-def _supports_cutlass(block_size: list[int] | None, output_dtype: torch.dtype) -> bool:
+def _supports_cutlass(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    block_size: list[int] | None,
+    output_dtype: torch.dtype,
+) -> bool:
     """
-    Check if CUTLASS blockwise FP8 matmul is supported for the given block size and output dtype.
+    Check if CUTLASS blockwise FP8 matmul is supported for the given inputs, output dtype, and block size.
 
     CUTLASS blockwise kernels require:
     - SM90+ (Hopper or newer)
     - Block size [128, 128] for weights
     - Block size [1, 128] for activations (handled implicitly)
     - Output dtype bfloat16 or float16
+    - K and N divisible by 16
     """
 
+    if torch.compiler.is_compiling():
+        # the checks after this, using importlib fail during torch.compile :/
+        return False
+
     if not is_torch_available() or not torch.cuda.is_available() or not is_kernels_available():
         return False
 
@@ -69,12 +103,17 @@ def _supports_cutlass(block_size: list[int] | None, output_dtype: torch.dtype) -
     if len(block_size) != 2 or block_size[0] != 128 or block_size[1] != 128:
         return False
 
+    # CUTLASS requires K and N divisible by 16
+    K, N = A.shape[-1], B.shape[0]
+    if K % 16 != 0 or N % 16 != 0:
+        return False
+
     # Check GPU capability (SM90+)
     capability = torch.cuda.get_device_capability()
     cuda_capability = capability[0] * 10 + capability[1]
 
     # Try to load the kernel and check if blockwise FP8 is supported
-    kernel = _get_quantization_kernel()
+    kernel = _get_cutlass_kernel()
     if kernel is None:
         return False
 
@@ -84,314 +123,46 @@ def _supports_cutlass(block_size: list[int] | None, output_dtype: torch.dtype) -
         return False
 
 
-_FP8_DTYPE = torch.float8_e4m3fn
-_FP8_MIN = torch.finfo(_FP8_DTYPE).min
-_FP8_MAX = torch.finfo(_FP8_DTYPE).max
-
-
-# Copied from https://huggingface.co/deepseek-ai/DeepSeek-V3/blob/main/inference/kernel.py
-@triton.jit
-def act_quant_kernel(x_ptr, y_ptr, s_ptr, BLOCK_SIZE: tl.constexpr):
-    pid = tl.program_id(axis=0)
-    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
-    x = tl.load(x_ptr + offs).to(tl.float32)
-    s = tl.max(tl.abs(x)) / 448.0
-    y = x / s
-    y = y.to(y_ptr.dtype.element_ty)
-    tl.store(y_ptr + offs, y)
-    tl.store(s_ptr + pid, s)
-
-
-def act_quant(x: torch.Tensor, block_size: int = 128) -> tuple[torch.Tensor, torch.Tensor]:
-    assert x.is_contiguous()
-    assert x.shape[-1] % block_size == 0
-    y = torch.empty_like(x, dtype=torch.float8_e4m3fn)
-    s = x.new_empty(*x.size()[:-1], x.size(-1) // block_size, dtype=torch.float32)
-
-    def grid(meta):
-        return (triton.cdiv(x.numel(), meta["BLOCK_SIZE"]),)
-
-    act_quant_kernel[grid](x, y, s, BLOCK_SIZE=block_size)
-    return y, s
-
-
-# Adapted from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/quantization/fp8_kernel.py
-@triton.jit
-def _w8a8_block_fp8_matmul(
-    # Pointers to inputs and output
-    A,
-    B,
-    C,
-    As,
-    Bs,
-    # Shape for matmul
-    M,
-    N,
-    K,
-    # Block size for block-wise quantization
-    group_n,
-    group_k,
-    # Stride for inputs and output
-    stride_am,
-    stride_ak,
-    stride_bk,
-    stride_bn,
-    stride_cm,
-    stride_cn,
-    stride_As_m,
-    stride_As_k,
-    stride_Bs_k,
-    stride_Bs_n,
-    # Meta-parameters
-    BLOCK_SIZE_M: tl.constexpr,
-    BLOCK_SIZE_N: tl.constexpr,
-    BLOCK_SIZE_K: tl.constexpr,
-    GROUP_SIZE_M: tl.constexpr,
-):
-    """Triton-accelerated function used to perform linear operations (dot
-    product) on input tensors `A` and `B` with block-wise quantization, and
-    store the result in output tensor `C`.
-    """
-
-    pid = tl.program_id(axis=0)
-    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
-    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
-    num_pid_in_group = GROUP_SIZE_M * num_pid_n
-    group_id = pid // num_pid_in_group
-    first_pid_m = group_id * GROUP_SIZE_M
-    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
-    pid_m = first_pid_m + (pid % group_size_m)
-    pid_n = (pid % num_pid_in_group) // group_size_m
-
-    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
-    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
-    offs_k = tl.arange(0, BLOCK_SIZE_K)
-    a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
-    b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
-
-    As_ptrs = As + offs_am * stride_As_m
-    offs_bsn = offs_bn // group_n
-    Bs_ptrs = Bs + offs_bsn * stride_Bs_n
-
-    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
-    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
-        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
-        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
-
-        k_start = k * BLOCK_SIZE_K
-        offs_ks = k_start // group_k
-        a_s = tl.load(As_ptrs + offs_ks * stride_As_k)
-        b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k)
-
-        accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
-        a_ptrs += BLOCK_SIZE_K * stride_ak
-        b_ptrs += BLOCK_SIZE_K * stride_bk
-
-    if C.dtype.element_ty == tl.bfloat16:
-        c = accumulator.to(tl.bfloat16)
-    elif C.dtype.element_ty == tl.float16:
-        c = accumulator.to(tl.float16)
-    else:
-        c = accumulator.to(tl.float32)
-
-    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
-    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-    c_ptrs = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
-    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
-    tl.store(c_ptrs, c, mask=c_mask)
-
-
-@triton.jit
-def _w8a8_block_fp8_matmul_per_tensor(
-    # Pointers to inputs and output
-    A,
-    B,
-    C,
-    As,
-    Bs,
-    # Shape for matmul
-    M,
-    N,
-    K,
-    # Block size for block-wise quantization
-    group_n,
-    group_k,
-    # Stride for inputs and output
-    stride_am,
-    stride_ak,
-    stride_bk,
-    stride_bn,
-    stride_cm,
-    stride_cn,
-    # Meta-parameters
-    BLOCK_SIZE_M: tl.constexpr,
-    BLOCK_SIZE_N: tl.constexpr,
-    BLOCK_SIZE_K: tl.constexpr,
-    GROUP_SIZE_M: tl.constexpr,
-):
-    """Triton-accelerated function used to perform linear operations (dot
-    product) on input tensors `A` and `B` with per-tensor quantization, and
-    store the result in output tensor `C`.
-    """
-
-    pid = tl.program_id(axis=0)
-    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
-    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
-    num_pid_in_group = GROUP_SIZE_M * num_pid_n
-    group_id = pid // num_pid_in_group
-    first_pid_m = group_id * GROUP_SIZE_M
-    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
-    pid_m = first_pid_m + (pid % group_size_m)
-    pid_n = (pid % num_pid_in_group) // group_size_m
-
-    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
-    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
-    offs_k = tl.arange(0, BLOCK_SIZE_K)
-    a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
-    b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
-    scale_a = tl.load(As)
-    scale_b = tl.load(Bs)
-    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
-    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
-        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
-        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
-
-        accumulator += tl.dot(a, b) * scale_a * scale_b
-        a_ptrs += BLOCK_SIZE_K * stride_ak
-        b_ptrs += BLOCK_SIZE_K * stride_bk
-
-    if C.dtype.element_ty == tl.bfloat16:
-        c = accumulator.to(tl.bfloat16)
-    elif C.dtype.element_ty == tl.float16:
-        c = accumulator.to(tl.float16)
-    else:
-        c = accumulator.to(tl.float32)
-
-    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
-    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-    c_ptrs = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
-    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
-    tl.store(c_ptrs, c, mask=c_mask)
-
-
-def w8a8_block_fp8_matmul_triton(
+def w8a8_block_fp8_matmul_cutlass(
     A: torch.Tensor,
     B: torch.Tensor,
     As: torch.Tensor,
     Bs: torch.Tensor,
-    block_size: list[int],
-    output_dtype: torch.dtype = torch.float32,
+    output_dtype: torch.dtype,
 ) -> torch.Tensor:
-    """This function performs matrix multiplication with block-wise
-    quantization.
-    It takes two input tensors `A` and `B` with scales `As` and `Bs`.
-    The output is returned in the specified `output_dtype`.
-    Args:
-        A: The input tensor, e.g., activation.
-        B: The input tensor, e.g., weight.
-        As: The per-token-group quantization scale for `A`.
-        Bs: The per-block quantization scale for `B`.
-        block_size: The block size for per-block quantization. It should
-        be 2-dim, e.g., [128, 128].
-        output_dytpe: The dtype of the returned tensor.
-    Returns:
-        torch.Tensor: The result of matmul.
-    """
-    if block_size is None:
-        block_n, block_k = 128, 128
-    else:
-        assert len(block_size) == 2
-        block_n, block_k = block_size[0], block_size[1]
+    """Call the CUTLASS blockwise FP8 matmul kernel.
 
-    # if we have per-tensor quantization, we use 128x128 block size for tiled matmul multiplication
-    if block_n == B.shape[-2] and block_k == B.shape[-1]:
-        block_n = 128
-        block_k = 128
+    Handles all layout conversions required by CUTLASS:
+      - A:  [M, K]           row-major    float8_e4m3fn
+      - B:  [K, N]           column-major float8_e4m3fn
+      - As: [M,  K//128]     M-major      (stride(0)==1)
+      - Bs: [K//128, N//128] K-major      (stride(0)==1)
+    """
+    kernel = _get_cutlass_kernel()
 
-    assert A.shape[-1] == B.shape[-1]
+    original_shape = A.shape
+    M = A.numel() // A.shape[-1]
+    K = A.shape[-1]
+    N = B.shape[0]
 
-    if As.numel() != 1:
-        assert A.shape[:-1] == As.shape[:-1] and A.is_contiguous()
-        assert triton.cdiv(A.shape[-1], block_k) == As.shape[-1]
+    A_2d = A.view(M, K).contiguous()
+    # B is [N, K] row-major; CUTLASS needs [K, N] column-major (stride(0)==1).
+    # .contiguous().t() gives [K, N] with stride=(1, K) — do NOT call .contiguous() after!
+    B_col_major = B.contiguous().t()
 
-    M = A.numel() // A.shape[-1]
+    # As: reshape to [M, K//128], then force M-major layout via t().contiguous().t()
+    As_2d = As.view(M, -1).contiguous()
+    As_2d = As_2d.t().contiguous().t()  # [M, K//128] with stride(0)==1
 
-    N, K = B.shape
-    assert B.ndim == 2 and B.is_contiguous()
-    if Bs.numel() != 1:
-        assert Bs.ndim == 2
-        assert triton.cdiv(N, block_n) == Bs.shape[0], f"{N}, {block_n}, {Bs.shape}"
-        assert triton.cdiv(K, block_k) == Bs.shape[1], f"{K}, {block_k}, {Bs.shape}"
-
-    C_shape = A.shape[:-1] + (N,)
-    C = A.new_empty(C_shape, dtype=output_dtype)
-
-    BLOCK_SIZE_M = 128
-    if M < BLOCK_SIZE_M:
-        BLOCK_SIZE_M = triton.next_power_of_2(M)
-        BLOCK_SIZE_M = max(BLOCK_SIZE_M, 16)
-    BLOCK_SIZE_K = block_k
-    assert block_k % BLOCK_SIZE_K == 0
-    BLOCK_SIZE_N = block_n
-
-    def grid(META):
-        return (triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),)
-
-    if As.numel() == 1 and Bs.numel() == 1:
-        _w8a8_block_fp8_matmul_per_tensor[grid](
-            A,
-            B,
-            C,
-            As,
-            Bs,
-            M,
-            N,
-            K,
-            block_n,
-            block_k,
-            A.stride(-2),
-            A.stride(-1),
-            B.stride(1),
-            B.stride(0),
-            C.stride(-2),
-            C.stride(-1),
-            BLOCK_SIZE_M=BLOCK_SIZE_M,
-            BLOCK_SIZE_N=BLOCK_SIZE_N,
-            BLOCK_SIZE_K=BLOCK_SIZE_K,
-            GROUP_SIZE_M=8,
-        )
-    else:
-        _w8a8_block_fp8_matmul[grid](
-            A,
-            B,
-            C,
-            As,
-            Bs,
-            M,
-            N,
-            K,
-            block_n,
-            block_k,
-            A.stride(-2),
-            A.stride(-1),
-            B.stride(1),
-            B.stride(0),
-            C.stride(-2),
-            C.stride(-1),
-            As.stride(-2),
-            As.stride(-1),
-            Bs.stride(1),
-            Bs.stride(0),
-            BLOCK_SIZE_M=BLOCK_SIZE_M,
-            BLOCK_SIZE_N=BLOCK_SIZE_N,
-            BLOCK_SIZE_K=BLOCK_SIZE_K,
-            GROUP_SIZE_M=8,
-        )
+    # Bs: our layout is [N//128, K//128]; CUTLASS needs [K//128, N//128] K-major (stride(0)==1)
+    Bs_km = Bs.contiguous().t()  # [K//128, N//128]
+    Bs_km = Bs_km.t().contiguous().t()  # force K-major (stride(0)==1)
 
-    return C
+    C = kernel.cutlass_scaled_mm(A_2d, B_col_major, As_2d, Bs_km, output_dtype, None)
+    return C.view(original_shape[:-1] + (N,))
 
 
-def w8a8_block_fp8_matmul(
+def w8a8_fp8_matmul(
     A: torch.Tensor,
     B: torch.Tensor,
     As: torch.Tensor,
@@ -412,146 +183,33 @@ def w8a8_block_fp8_matmul(
     Otherwise falls back to Triton.
     """
 
-    if _supports_cutlass(block_size, output_dtype):
-        kernel = _get_quantization_kernel()
-        if kernel is not None:
-            try:
-                # CUTLASS expects:
-                # - A: [M, K] row-major, float8_e4m3fn
-                # - B: [K, N] column-major, float8_e4m3fn
-                # - As: [M, K//128] M-major (activation scales)
-                # - Bs: [K//128, N//128] K-major (weight scales)
-
-                # Reshape A to 2D if needed
-                original_shape = A.shape
-                M = A.numel() // A.shape[-1]
-                K = A.shape[-1]
-                N = B.shape[0]
-
-                # CUTLASS requires dimensions divisible by 16
-                if K % 16 != 0 or N % 16 != 0:
-                    raise ValueError(f"CUTLASS requires K ({K}) and N ({N}) divisible by 16")
-
-                A_2d = A.view(M, K).contiguous()
-                # B needs to be column-major for CUTLASS: [K, N] with stride(0)==1
-                # Our B is [N, K] row-major. Make it contiguous first, then transpose.
-                # B.contiguous() gives [N, K] with stride=(K,1)
-                # B.contiguous().t() gives [K, N] with stride=(1,K) which is column-major
-                # Do NOT call .contiguous() after .t() as it would make it row-major!
-                B_col_major = B.contiguous().t()
-
-                # Scales need proper layout for CUTLASS blockwise:
-                # As should be [M, K//128] with M-major layout (stride(0)==1)
-                # Bs should be [K//128, N//128] with K-major layout (stride(0)==1)
-
-                # As: reshape to [M, K//128], then make M-major via t().contiguous().t()
-                As_2d = As.view(M, -1).contiguous()
-                As_2d = As_2d.t().contiguous().t()  # [M, K//128] with stride(0)==1
-
-                # Bs: our input is [N//128, K//128], need [K//128, N//128] with stride(0)==1
-                # Transpose to get [K//128, N//128], then make K-major via t().contiguous().t()
-                Bs_km = Bs.contiguous().t()  # [K//128, N//128]
-                Bs_km = Bs_km.t().contiguous().t()  # Make K-major (stride(0)==1)
-
-                # Call CUTLASS kernel - it returns the output tensor
-                # Signature: cutlass_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias=None) -> Tensor
-                C = kernel.cutlass_scaled_mm(A_2d, B_col_major, As_2d, Bs_km, output_dtype, None)
-                # Reshape output back
-                C_shape = original_shape[:-1] + (N,)
-                return C.view(C_shape)
-            except Exception as e:
-                logger.warning_once(f"CUTLASS kernel failed: {e}. Falling back to Triton.")
+    if _supports_cutlass(A, B, block_size, output_dtype):
+        return w8a8_block_fp8_matmul_cutlass(A, B, As, Bs, output_dtype)
 
     # Fall back to Triton
-    return w8a8_block_fp8_matmul_triton(A, B, As, Bs, block_size, output_dtype)
-
-
-# Python version of the above triton function, it's much slower than the triton version, for testing
-@torch.compile
-def w8a8_block_fp8_matmul_compile(
-    input_q: torch.Tensor,  # [batch, seq_len, hidden_dim]
-    weight_q: torch.Tensor,  # [out_features, hidden_dim]
-    input_scale: torch.Tensor,  # [batch * seq_len, num_input_groups]
-    weight_scale: torch.Tensor,  # [num_weight_blocks_m, num_weight_blocks_n]
-    block_size: tuple[int, int] | None = None,  # (M=128, N=128) for weights for example
-    output_dtype: torch.dtype = torch.float32,
-) -> torch.Tensor:
-    """
-    Performs blocked matrix multiplication with FP8 quantized matrices.
-
-    Args:
-        input_q: Quantized input tensor with 1x128 block quantization
-        weight_q: Quantized weight tensor with 128x128 block quantization
-        input_scale: Scaling factors for input blocks
-        weight_scale: Scaling factors for weight blocks
-        block_size: Tuple of (M, N) for weight block dimensions
-        output_dtype: Desired output dtype
-    """
-    batch_size, seq_len, hidden_dim = input_q.shape if input_q.ndim == 3 else (1, input_q.shape[0], input_q.shape[1])
-    out_features = weight_q.shape[0]
-
-    # Reshape input for batched matmul
-    input_reshaped = input_q.view(-1, hidden_dim)  # [batch*seq_len, hidden_dim]
-    input_scale_reshaped = input_scale.view(input_scale.shape[0], -1)  # [batch*seq_len, 1]
-    # Calculate number of blocks
-    num_weight_blocks_m = out_features // block_size[0]
-    num_weight_blocks_n = hidden_dim // block_size[1]
-
-    output = torch.zeros((batch_size * seq_len, out_features), dtype=torch.float32, device=input_q.device)
-
-    for i in range(num_weight_blocks_m):
-        m_start = i * block_size[0]
-        m_end = m_start + block_size[0]
-
-        for j in range(num_weight_blocks_n):
-            n_start = j * block_size[1]
-            n_end = n_start + block_size[1]
-
-            # Extract current blocks
-            input_block = input_reshaped[:, n_start:n_end]
-            weight_block = weight_q[m_start:m_end, n_start:n_end]
-
-            # Get corresponding scales
-            curr_input_scale = input_scale_reshaped[:, j : j + 1]  # [batch*seq_len, 1]
-            curr_weight_scale = weight_scale[i, j]  # scalar
-
-            block_result = (
-                torch._scaled_mm(
-                    input_block,
-                    weight_block.t(),
-                    scale_a=torch.tensor(1, dtype=torch.float32, device=input_q.device),
-                    scale_b=curr_weight_scale,
-                    out_dtype=output_dtype,
-                )
-                * curr_input_scale
-            )
-
-            output[:, m_start:m_end] += block_result
-
-    output = output.view(batch_size, seq_len, out_features)
-
-    return output.to(output_dtype)
+    kernel = _get_triton_kernel()
+    return kernel.w8a8_fp8_matmul(A, B, As, Bs, block_size, output_dtype)
 
 
 class FP8Linear(nn.Linear):
     def __init__(
         self,
         in_features: int,
         out_features: int,
-        bias: bool = False,
-        dtype=torch.float8_e4m3fn,
         block_size: tuple[int, int] | None = None,
-        activation_scheme="dynamic",
+        activation_scheme: str = "dynamic",
+        has_bias: bool = False,
+        dtype=_FP8_DTYPE,
     ):
         super().__init__(in_features, out_features)
 
-        # If block size is None, it means that we are doing per-tensor quantization
+        self.has_bias = has_bias
         self.block_size = block_size
         self.activation_scheme = activation_scheme
-
         self.weight = torch.nn.Parameter(torch.empty(out_features, in_features, dtype=dtype))
 
         if self.block_size is None:
+            # If block size is None, it means that we are doing per-tensor quantization
             self.weight_scale_inv = nn.Parameter(torch.tensor(1.0, dtype=torch.float32))
         else:
             scale_out_features = (out_features + self.block_size[0] - 1) // self.block_size[0]
@@ -562,157 +220,308 @@ def __init__(
 
         if self.activation_scheme == "static":
             self.activation_scale = nn.Parameter(torch.tensor(1.0, dtype=torch.float32))
+        else:
+            self.register_parameter("activation_scale", None)
 
-        if bias:
+        if self.has_bias:
             self.bias = nn.Parameter(torch.empty(self.out_features))
         else:
             self.register_parameter("bias", None)
 
     def forward(self, input: torch.Tensor) -> torch.Tensor:
         if self.weight.element_size() > 1:
             return F.linear(input, self.weight, self.bias)
+
+        if isinstance(self.weight, torch.distributed.tensor.DTensor):
+            weight = self.weight._local_tensor.contiguous()
+            scale_inv = self.weight_scale_inv._local_tensor.contiguous()
         else:
-            if isinstance(self.weight, torch.distributed.tensor.DTensor):
-                weight = self.weight._local_tensor.contiguous()
-                scale_inv = self.weight_scale_inv._local_tensor.contiguous()
-            else:
-                weight = self.weight.contiguous()
-                scale_inv = self.weight_scale_inv.contiguous()
-            # Context manager used to switch among the available accelerators
-            device_type = torch.accelerator.current_accelerator().type if is_torch_accelerator_available() else "cuda"
-            torch_accelerator_module = getattr(torch, device_type, torch.cuda)
-            with torch_accelerator_module.device(input.device):
-                if self.activation_scheme == "dynamic":
-                    qinput, scale = act_quant(input, self.block_size[1])
-                elif self.activation_scheme == "static":
-                    scale = self.activation_scale.to(torch.float32)
-                    qinput = (input / scale).clamp(min=_FP8_MIN, max=_FP8_MAX).to(torch.float8_e4m3fn)
-
-                else:
-                    raise NotImplementedError("Not supported")
-
-                output = w8a8_block_fp8_matmul(
-                    qinput,
-                    weight,
-                    scale,
-                    scale_inv,
-                    self.block_size,
-                    output_dtype=input.dtype,
-                )
+            # why wouldn't it be contiguous?
+            weight = self.weight.contiguous()
+            scale_inv = self.weight_scale_inv.contiguous()
+
+        if self.activation_scheme == "dynamic":
+            kernel = _get_triton_kernel()
+            qinput, scale = kernel.fp8_act_quant(
+                input, self.block_size[1] if self.block_size is not None else input.shape[-1]
+            )
+        elif self.activation_scheme == "static":
+            scale = self.activation_scale.to(torch.float32)
+            qinput = (input / scale).clamp(min=_FP8_MIN, max=_FP8_MAX).to(_FP8_DTYPE)
+        else:
+            raise NotImplementedError(f"Unsupported activation scheme: {self.activation_scheme}")
+
+        output = w8a8_fp8_matmul(
+            qinput,
+            weight,
+            scale,
+            scale_inv,
+            self.block_size,
+            output_dtype=input.dtype,
+        )
+
+        if self.bias is not None:
+            output = output + self.bias
+
+        return output.to(dtype=input.dtype)
+
+
+def fp8_batched_mm_experts_forward(
+    self: torch.nn.Module,
+    hidden_states: torch.Tensor,
+    top_k_index: torch.Tensor,
+    top_k_weights: torch.Tensor,
+) -> torch.Tensor:
+    kernel = _get_triton_kernel()
+    device = hidden_states.device
+    num_top_k = top_k_index.size(-1)
+    num_tokens = hidden_states.size(0)
+    hidden_dim = hidden_states.size(-1)
+
+    # Reshape for easier indexing
+    # S is the number of selected tokens-experts pairs (S = num_tokens * num_top_k)
+    token_idx = torch.arange(num_tokens, device=device).unsqueeze(1).expand(-1, num_top_k).reshape(-1)  # (S,)
+    sample_weights = top_k_weights.reshape(-1)  # (S,)
+    expert_ids = top_k_index.reshape(-1)  # (S,)
+
+    # Get current hidden states for selected samples
+    selected_hidden_states = hidden_states[token_idx]
+
+    # --- Up projection per expert (FP8 batched) ---
+    proj_out = kernel.w8a8_fp8_matmul_batched(
+        selected_hidden_states,
+        self.gate_up_proj if self.has_gate else self.up_proj,
+        self.gate_up_proj_scale_inv if self.has_gate else self.up_proj_scale_inv,
+        block_size=self.block_size,
+        expert_ids=expert_ids,
+    )  # (S, 2 * intermediate_dim) or (S, intermediate_dim) depending on gating
+
+    # Apply gating or activation
+    if self.has_gate:
+        # for gated experts we apply the custom/default gating mechanism
+        proj_out = self._apply_gate(proj_out)  # (S, intermediate_dim)
+    else:
+        # for non-gated experts we just apply the activation function
+        proj_out = self.act_fn(proj_out)  # (S, intermediate_dim)
+
+    # --- Down projection per expert (FP8 batched) ---
+    proj_out = kernel.w8a8_fp8_matmul_batched(
+        proj_out,
+        self.down_proj,
+        self.down_proj_scale_inv,
+        block_size=self.block_size,
+        expert_ids=expert_ids,
+    )  # (S, hidden_dim)
+
+    # Apply routing weights
+    weighted_out = proj_out * sample_weights.to(proj_out.dtype).unsqueeze(-1)  # (S, hidden_dim)
+
+    # Accumulate results using deterministic reshape+sum instead of index_add_
+    # (index_add_ with duplicate indices is non-deterministic on CUDA due to atomicAdd)
+    final_hidden_states = weighted_out.view(num_tokens, num_top_k, hidden_dim).sum(dim=1)
 
-            # Blocks the CPU until all accelerator operations on the specified device are complete. It is used to ensure that the results of the
-            # preceding operations are ready before proceeding
-            torch_accelerator_module.synchronize()
-            if self.bias is not None:
-                output = output + self.bias
+    return final_hidden_states.to(hidden_states.dtype)
 
-            return output.to(dtype=input.dtype)
 
+def fp8_grouped_mm_experts_forward(
+    self: torch.nn.Module,
+    hidden_states: torch.Tensor,
+    top_k_index: torch.Tensor,
+    top_k_weights: torch.Tensor,
+) -> torch.Tensor:
+    kernel = _get_triton_kernel()
+    device = hidden_states.device
+    num_top_k = top_k_index.size(-1)
+    num_tokens = hidden_states.size(0)
+    hidden_dim = hidden_states.size(-1)
+
+    # S is the number of selected token-expert pairs (S = num_tokens * num_top_k)
+    token_idx = torch.arange(num_tokens, device=device).unsqueeze(1).expand(-1, num_top_k).reshape(-1)  # (S,)
+    sample_weights = top_k_weights.reshape(-1)  # (S,)
+    expert_ids = top_k_index.reshape(-1)  # (S,)
+
+    selected_hidden_states = hidden_states[token_idx]
+
+    # Sort by expert for grouped processing
+    perm = torch.argsort(expert_ids)
+    inv_perm = torch.empty_like(perm)
+    inv_perm[perm] = torch.arange(perm.size(0), device=device)
+
+    expert_ids_g = expert_ids[perm]
+    sample_weights_g = sample_weights[perm]
+    selected_hidden_states_g = selected_hidden_states[perm]
+
+    # Compute offsets for grouped processing.
+    # histc instead of bincount avoids cuda-graph issues;
+    # CPU requires float input, CUDA requires int input (deterministic mode).
+    histc_input = expert_ids_g.float() if device.type == "cpu" else expert_ids_g.int()
+    tokens_per_expert = torch.histc(histc_input, bins=self.num_experts, min=0, max=self.num_experts - 1)
+    offsets = torch.cumsum(tokens_per_expert, dim=0, dtype=torch.int32)
+
+    # --- Up projection per expert (FP8 grouped) ---
+    proj_out = kernel.w8a8_fp8_matmul_grouped(
+        selected_hidden_states_g,
+        self.gate_up_proj if self.has_gate else self.up_proj,
+        self.gate_up_proj_scale_inv if self.has_gate else self.up_proj_scale_inv,
+        tokens_per_expert=tokens_per_expert,
+        block_size=self.block_size,
+        offsets=offsets,
+    )  # (S, 2 * intermediate_dim)
+
+    # Apply gating or activation
+    if self.has_gate:
+        # for gated experts we apply the custom/default gating mechanism
+        proj_out = self._apply_gate(proj_out)  # (S, intermediate_dim)
+    else:
+        # for non-gated experts we just apply the activation function
+        proj_out = self.act_fn(proj_out)  # (S, intermediate_dim)
+
+    # --- Down projection per expert (FP8 grouped) ---
+    proj_out = kernel.w8a8_fp8_matmul_grouped(
+        proj_out,
+        self.down_proj,
+        self.down_proj_scale_inv,
+        tokens_per_expert=tokens_per_expert,
+        block_size=self.block_size,
+        offsets=offsets,
+    )  # (S, hidden_dim)
+
+    # Apply routing weights
+    weighted_out = proj_out * sample_weights_g.to(proj_out.dtype).unsqueeze(-1)  # (S, hidden_dim)
 
-def _ceil_div(a, b):
-    return (a + b - 1) // b
+    # Restore original order
+    weighted_out = weighted_out[inv_perm]
 
+    # Accumulate results using deterministic reshape+sum instead of index_add_
+    # (index_add_ with duplicate indices is non-deterministic on CUDA due to atomicAdd)
+    final_hidden_states = weighted_out.view(num_tokens, num_top_k, hidden_dim).sum(dim=1)
 
-class FP8Expert(nn.Module):
-    def __init__(self, config, block_size, dtype=torch.float8_e4m3fn):
+    return final_hidden_states.to(hidden_states.dtype)
+
+
+class FP8Experts(nn.Module):
+    def __init__(
+        self,
+        config,
+        block_size: tuple[int, int] | None = None,
+        activation_scheme: str = "dynamic",
+        has_bias: bool = False,
+        has_gate: bool = True,
+        dtype=_FP8_DTYPE,
+    ):
         super().__init__()
 
-        from ..activations import ACT2FN
+        assert has_bias is False, (
+            "FP8Experts does not support bias for now, please open an issue if you want this feature"
+        )
+        assert activation_scheme == "dynamic", (
+            "Only dynamic activation quantization is supported for now, please open an issue if you want others"
+        )
 
+        self.config = config
+        self.has_bias = has_bias
+        self.has_gate = has_gate
         self.block_size = block_size
-        # TODO we don't need exact expert count here but only in forward
-        self.num_experts = config.num_local_experts if hasattr(config, "num_local_experts") else config.num_experts
         self.hidden_dim = config.hidden_size
+        self.activation_scheme = activation_scheme
+        self.num_experts = config.num_local_experts if hasattr(config, "num_local_experts") else config.num_experts
         self.intermediate_dim = (
             config.moe_intermediate_size if hasattr(config, "moe_intermediate_size") else config.intermediate_size
         )
 
-        Wg_out, Wg_in = 2 * self.intermediate_dim, self.hidden_dim
-        Wd_out, Wd_in = self.hidden_dim, self.intermediate_dim
-
-        self.gate_up_proj = nn.Parameter(torch.zeros(self.num_experts, Wg_out, Wg_in, dtype=dtype))
-        self.down_proj = nn.Parameter(torch.zeros(self.num_experts, Wd_out, Wd_in, dtype=dtype))
-
-        bo, bi = self.block_size
-
-        # gate_up tiles: ceil(Wg_out/bo) x ceil(Wg_in/bi)
-        gu_scale_o = _ceil_div(Wg_out, bo)
-        gu_scale_i = _ceil_div(Wg_in, bi)
-        self.gate_up_proj_scale_inv = nn.Parameter(
-            torch.zeros(self.num_experts, gu_scale_o, gu_scale_i, dtype=torch.float32)
-        )
+        if self.has_gate:
+            gu_proj_out, gu_proj_in = 2 * self.intermediate_dim, self.hidden_dim
+            self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, gu_proj_out, gu_proj_in, dtype=dtype))
+            gu_scale_out = triton.cdiv(gu_proj_out, self.block_size[0]) if self.block_size is not None else 1
+            gu_scale_in = triton.cdiv(gu_proj_in, self.block_size[1]) if self.block_size is not None else 1
+            self.gate_up_proj_scale_inv = nn.Parameter(
+                torch.empty(self.num_experts, gu_scale_out, gu_scale_in, dtype=torch.float32)
+            )
+            self.register_parameter("gate_up_proj_bias", None)
+        else:
+            u_proj_out, u_proj_in = self.intermediate_dim, self.hidden_dim
+            self.up_proj = nn.Parameter(torch.empty(self.num_experts, u_proj_out, u_proj_in, dtype=dtype))
+            u_scale_out = triton.cdiv(u_proj_out, self.block_size[0]) if self.block_size is not None else 1
+            u_scale_in = triton.cdiv(u_proj_in, self.block_size[1]) if self.block_size is not None else 1
+            self.up_proj_scale_inv = nn.Parameter(
+                torch.empty(self.num_experts, u_scale_out, u_scale_in, dtype=torch.float32)
+            )
+            self.register_parameter("up_proj_bias", None)
 
-        # down tiles: ceil(Wd_out/bo) x ceil(Wd_in/bi)
-        dp_scale_o = _ceil_div(Wd_out, bo)
-        dp_scale_i = _ceil_div(Wd_in, bi)
+        d_proj_out, d_proj_in = self.hidden_dim, self.intermediate_dim
+        self.down_proj = nn.Parameter(torch.empty(self.num_experts, d_proj_out, d_proj_in, dtype=dtype))
+        d_scale_out = triton.cdiv(d_proj_out, self.block_size[0]) if self.block_size is not None else 1
+        d_scale_in = triton.cdiv(d_proj_in, self.block_size[1]) if self.block_size is not None else 1
         self.down_proj_scale_inv = nn.Parameter(
-            torch.zeros(self.num_experts, dp_scale_o, dp_scale_i, dtype=torch.float32)
+            torch.empty(self.num_experts, d_scale_out, d_scale_in, dtype=torch.float32)
         )
+        self.register_parameter("down_proj_bias", None)
 
-        # (Optional) bias per projection — many MoEs omit bias; keep None to match your FP8Linear default
-        self.register_parameter("gate_up_bias", None)
-        self.register_parameter("down_bias", None)
-
-        # Activation used in the MLP (same as your config / ACT2FN)
-        # Keep a handle here; actual usage happens in forward of your MoE block
         self.act_fn = ACT2FN[config.hidden_act]
 
-    # We follow the mixtral "eager" moe implementation at
-    # https://github.com/huggingface/transformers/blob/457048fbfdba9a7dee8bd03328c62f49e57b95f9/src/transformers/models/mixtral/modular_mixtral.py#L148
-    # The core changes in this FP8 version should only relate to how we call the linear projections
+    def _apply_gate(self, gate_up: torch.Tensor) -> torch.Tensor:
+        gate, up = gate_up.chunk(2, dim=-1)
+        return self.act_fn(gate) * up
+
     def forward(
-        self,
-        hidden_states: torch.Tensor,
-        top_k_index: torch.Tensor,
-        top_k_weights: torch.Tensor,
+        self, hidden_states: torch.Tensor, top_k_index: torch.Tensor, top_k_weights: torch.Tensor
     ) -> torch.Tensor:
-        final_hidden_states = torch.zeros_like(hidden_states)
+        # index_add_ will accumulate using the dtype of the tensor we write into
+        # so we use float32 for the accumulation to avoid numerical issues in bf16/fp16
+        final_hidden_states = torch.zeros_like(hidden_states, dtype=torch.float32)
+
         with torch.no_grad():
             expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts)
             expert_mask = expert_mask.permute(2, 1, 0)
-            expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
+            expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero(as_tuple=False).view(-1)
 
         for expert_idx in expert_hit:
-            expert_idx = expert_idx[0]
-            if expert_idx == len(self.gate_up_proj):  # weights will load fine
+            if expert_idx == self.num_experts:
                 continue
+
             top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
             current_state = hidden_states[token_idx]
-            gate, up = self.linear(
-                current_state, self.gate_up_proj[expert_idx], self.gate_up_proj_scale_inv[expert_idx]
-            ).chunk(2, dim=-1)
-            current_hidden_states = self.act_fn(gate) * up
-            current_hidden_states = self.linear(
-                current_hidden_states, self.down_proj[expert_idx], self.down_proj_scale_inv[expert_idx]
+            proj_out = self.linear(
+                current_state,
+                self.gate_up_proj[expert_idx] if self.has_gate else self.up_proj[expert_idx],
+                self.gate_up_proj_scale_inv[expert_idx] if self.has_gate else self.up_proj_scale_inv[expert_idx],
             )
-
+            proj_out = self._apply_gate(proj_out) if self.has_gate else self.act_fn(proj_out)
+            proj_out = self.linear(proj_out, self.down_proj[expert_idx], self.down_proj_scale_inv[expert_idx])
             routing_weights = top_k_weights[token_idx, top_k_pos, None]
-            current_hidden_states = current_hidden_states * routing_weights.to(current_hidden_states.dtype)
-            final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype))
-
-        return final_hidden_states
+            weighted_out = proj_out * routing_weights.to(proj_out.dtype)
+            final_hidden_states.index_add_(0, token_idx, weighted_out.to(final_hidden_states.dtype))
+        return final_hidden_states.to(hidden_states.dtype)
 
     def linear(self, input: torch.Tensor, weight: torch.Tensor, weight_scale_inv: torch.Tensor) -> torch.Tensor:
         if weight.element_size() > 1:
             return F.linear(input, weight, None)
-        else:
-            # Context manager used to switch among the available accelerators
-            device_type = torch.accelerator.current_accelerator().type if is_torch_accelerator_available() else "cuda"
-            torch_accelerator_module = getattr(torch, device_type, torch.cuda)
-            with torch_accelerator_module.device(input.device):
-                qinput, scale = act_quant(input, self.block_size[1])
-                output = w8a8_block_fp8_matmul(
-                    qinput,
-                    weight,
-                    scale,
-                    weight_scale_inv,
-                    self.block_size,
-                    output_dtype=input.dtype,
-                )
-            # Blocks the CPU until all accelerator operations on the specified device are complete. It is used to ensure that the results of the
-            # preceding operations are ready before proceeding
-            torch_accelerator_module.synchronize()
-            return output.to(dtype=input.dtype)
+
+        kernel = _get_triton_kernel()
+        qinput, scale = kernel.fp8_act_quant(
+            input, self.block_size[1] if self.block_size is not None else input.shape[-1]
+        )
+        output = w8a8_fp8_matmul(
+            qinput,
+            weight,
+            scale,
+            weight_scale_inv,
+            self.block_size,
+            output_dtype=input.dtype,
+        )
+        return output.to(dtype=input.dtype)
+
+
+class FP8ExpertsInterface(ExpertsInterface):
+    """Interface for registering custom FP8 experts forward functions."""
+
+    _global_mapping = {
+        "batched_mm": fp8_batched_mm_experts_forward,
+        "grouped_mm": fp8_grouped_mm_experts_forward,
+    }
+
+
+ALL_FP8_EXPERTS_FUNCTIONS = FP8ExpertsInterface()
 
 
 def replace_with_fp8_linear(
@@ -739,23 +548,36 @@ def replace_with_fp8_linear(
     for module_name, module in model.named_modules():
         if not should_convert_module(module_name, modules_to_not_convert):
             continue
+
         # we need this to correctly materialize the weights during quantization
         module_kwargs = {} if pre_quantized else {"dtype": None}
         new_module = None
         with torch.device("meta"):
             if module_name.endswith(".experts"):
-                new_module = FP8Expert(
-                    config=model.config.get_text_config(),
+                has_gate = getattr(module, "has_gate", True)
+                has_bias = getattr(module, "has_bias", False)
+                config = getattr(module, "config", model.config.get_text_config())
+                new_class = use_experts_implementation(
+                    experts_class=FP8Experts,
+                    experts_interface=ALL_FP8_EXPERTS_FUNCTIONS,
+                    has_bias=has_bias,
+                    has_gate=has_gate,
+                )
+                new_module = new_class(
+                    config=config,
                     block_size=quantization_config.weight_block_size,
+                    activation_scheme=quantization_config.activation_scheme,
+                    has_bias=has_bias,
+                    has_gate=has_gate,
                     **module_kwargs,
                 )
             elif isinstance(module, nn.Linear):
                 new_module = FP8Linear(
                     in_features=module.in_features,
                     out_features=module.out_features,
-                    bias=module.bias is not None,
-                    activation_scheme=quantization_config.activation_scheme,
                     block_size=quantization_config.weight_block_size,
+                    activation_scheme=quantization_config.activation_scheme,
+                    has_bias=module.bias is not None,
                     **module_kwargs,
                 )
             if new_module is not None:

filename: src/transformers/integrations/moe.py
patch: @@ -17,25 +17,26 @@
 
 from ..utils import logging
 from ..utils.generic import GeneralInterface
-from ..utils.import_utils import is_grouped_mm_available, is_torch_available, is_torchdynamo_compiling
+from ..utils.import_utils import is_torch_available, is_torchdynamo_compiling
 
 
 if is_torch_available():
     import torch
 
+
 logger = logging.get_logger(__name__)
 
 # Examples of experts class with its eager mm implementation
-# class Experts(nn.Module):
+# class Experts(torch.nn.Module):
 #     """Collection of expert weights stored as 3D tensors."""
 
 #     def __init__(self, config):
 #         super().__init__()
 #         self.num_experts = config.n_routed_experts
 #         self.hidden_dim = config.hidden_size
 #         self.intermediate_dim = config.moe_intermediate_size
-#         self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim))
-#         self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim))
+#         self.gate_up_proj = torch.nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim))
+#         self.down_proj = torch.nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim))
 #         self.act_fn = ACT2FN[config.hidden_act]
 
 #     def forward(
@@ -56,9 +57,9 @@
 #                 continue
 #             top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
 #             current_state = hidden_states[token_idx]
-#             gate, up = nn.functional.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1)
+#             gate, up = torch.nn.functional.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1)
 #             current_hidden_states = self.act_fn(gate) * up
-#             current_hidden_states = nn.functional.linear(current_hidden_states, self.down_proj[expert_idx])
+#             current_hidden_states = torch.nn.functional.linear(current_hidden_states, self.down_proj[expert_idx])
 #             current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None]
 #             final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype))
 
@@ -133,28 +134,29 @@ def batched_mm_experts_forward(
         selected_biases = self.up_proj_bias[expert_ids] if self.has_bias else None
 
     # --- Up projection per expert (batched) ---
-    up_proj_out = _batched_linear(
+    proj_out = _batched_linear(
         selected_hidden_states, selected_weights, bias=selected_biases, is_transposed=self.is_transposed
     )  # (S, 2 * intermediate_dim) or  (S, intermediate_dim) depending on whether we have gating
 
-    # Apply gating or just activation
+    # Apply gating or activation
     if self.has_gate:
-        up_proj_out = self._apply_gate(up_proj_out)  # (S, intermediate_dim)
+        # for gated experts we apply the custom/default gating mechanism
+        proj_out = self._apply_gate(proj_out)  # (S, intermediate_dim)
     else:
         # for non-gated experts we just apply the activation function
-        up_proj_out = self.act_fn(up_proj_out)  # (S, intermediate_dim)
+        proj_out = self.act_fn(proj_out)  # (S, intermediate_dim)
 
     # Select down projection weights and biases for selected samples
     selected_weights = self.down_proj[expert_ids]
     selected_biases = self.down_proj_bias[expert_ids] if self.has_bias else None
 
     # --- Down projection per expert (batched) ---
-    down_proj_out = _batched_linear(
-        up_proj_out, selected_weights, bias=selected_biases, is_transposed=self.is_transposed
+    proj_out = _batched_linear(
+        proj_out, selected_weights, bias=selected_biases, is_transposed=self.is_transposed
     )  # (S, hidden_dim)
 
     # Apply routing weights and zero out invalid expert contributions
-    weighted_out = down_proj_out * sample_weights.unsqueeze(-1)  # (S, hidden_dim)
+    weighted_out = proj_out * sample_weights.unsqueeze(-1)  # (S, hidden_dim)
     weighted_out.masked_fill_(invalid_mask.unsqueeze(-1), 0.0)  # Zero out invalid expert contributions
 
     # Accumulate results using deterministic reshape+sum instead of index_add_
@@ -263,7 +265,7 @@ def _can_use_grouped_mm(input: torch.Tensor, weight: torch.Tensor, offs: torch.T
         # torch.grouped_mm is not supported in torch.compile with dtypes other than bfloat16
         return False
 
-    return is_grouped_mm_available()
+    return hasattr(torch.nn.functional, "grouped_mm") or hasattr(torch, "_grouped_mm")
 
 
 def _grouped_mm(
@@ -357,7 +359,9 @@ def grouped_mm_experts_forward(
 
     # Sort by expert for grouped processing
     perm = torch.argsort(expert_ids)
-    inv_perm = torch.argsort(perm)
+    inv_perm = torch.empty_like(perm)
+    inv_perm[perm] = torch.arange(perm.size(0), device=device)
+
     expert_ids_g = expert_ids[perm]
     sample_weights_g = sample_weights[perm]
     selected_hidden_states_g = selected_hidden_states[perm]
@@ -366,8 +370,8 @@ def grouped_mm_experts_forward(
     # using histc instead of bincount to avoid cuda graph issues
     # With deterministic algorithms, CPU only supports float input, CUDA only supports int input.
     histc_input = expert_ids_g.float() if device.type == "cpu" else expert_ids_g.int()
-    num_tokens_per_expert = torch.histc(histc_input, bins=self.num_experts, min=0, max=self.num_experts - 1)
-    offsets = torch.cumsum(num_tokens_per_expert, dim=0, dtype=torch.int32)
+    tokens_per_expert = torch.histc(histc_input, bins=self.num_experts, min=0, max=self.num_experts - 1)
+    offsets = torch.cumsum(tokens_per_expert, dim=0, dtype=torch.int32)
 
     # Select expert weights and biases
     # NOTE: We keep all experts here and rely on offsets to target the active ones.
@@ -382,28 +386,29 @@ def grouped_mm_experts_forward(
         selected_biases = self.up_proj_bias[expert_ids_g] if self.has_bias else None
 
     # --- Up projection per expert (grouped) ---
-    up_proj_out = _grouped_linear(
+    proj_out = _grouped_linear(
         selected_hidden_states_g, selected_weights, offsets, bias=selected_biases, is_transposed=self.is_transposed
     )  # (S, 2 * intermediate_dim) or  (S, intermediate_dim) depending on whether we have gating
 
-    # Apply gating or just activation
+    # Apply gating or activation
     if self.has_gate:
-        up_proj_out = self._apply_gate(up_proj_out)  # (S, intermediate_dim)
+        # for gated experts we apply the custom/default gating mechanism
+        proj_out = self._apply_gate(proj_out)  # (S, intermediate_dim)
     else:
         # for non-gated experts we just apply the activation function
-        up_proj_out = self.act_fn(up_proj_out)  # (S, intermediate_dim)
+        proj_out = self.act_fn(proj_out)  # (S, intermediate_dim)
 
     # Select down projection weights and biases
     selected_weights = self.down_proj
     selected_biases = self.down_proj_bias[expert_ids_g] if self.has_bias else None
 
     # --- Down projection per expert (grouped) ---
-    down_proj_out = _grouped_linear(
-        up_proj_out, selected_weights, offsets, bias=selected_biases, is_transposed=self.is_transposed
+    proj_out = _grouped_linear(
+        proj_out, selected_weights, offsets, bias=selected_biases, is_transposed=self.is_transposed
     )  # (S, hidden_dim)
 
     # Apply routing weights
-    weighted_out = down_proj_out * sample_weights_g.unsqueeze(-1)  # (S, hidden_dim)
+    weighted_out = proj_out * sample_weights_g.unsqueeze(-1)  # (S, hidden_dim)
 
     # Restore original order
     weighted_out = weighted_out[inv_perm]  # (S, hidden_dim)
@@ -418,7 +423,7 @@ def grouped_mm_experts_forward(
 
 
 class ExpertsInterface(GeneralInterface):
-    """Interface for registering custom experts implementations."""
+    """Interface for registering custom experts forward functions."""
 
     _global_mapping = {
         "batched_mm": batched_mm_experts_forward,
@@ -460,6 +465,7 @@ def _default_apply_gate(self, gate_up_out: torch.Tensor) -> torch.Tensor:
 def use_experts_implementation(
     experts_class: type[torch.nn.Module] | None = None,
     *,
+    experts_interface: ExpertsInterface = ALL_EXPERTS_FUNCTIONS,
     is_transposed: bool = False,
     has_bias: bool = False,
     has_gate: bool = True,
@@ -469,6 +475,8 @@ def use_experts_implementation(
     Args:
         experts_class (`type[torch.nn.Module]`, *optional*):
             The experts class to modify. If not provided, returns a decorator that can be applied to the class.
+        experts_interface (`ExpertsInterface`, *optional*, defaults to `ALL_EXPERTS_FUNCTIONS`):
+            The experts interface to use for dispatching the forward method.
         is_transposed (`bool`, *optional*, defaults to `False`):
             Whether the expert weights are stored in transposed format.
         has_bias (`bool`, *optional*, defaults to `False`):
@@ -492,9 +500,7 @@ def __init__(self, config, *args, **kwargs):
 
         @wraps(original_forward)
         def forward(self, *args, **kwargs):
-            experts_forward = ALL_EXPERTS_FUNCTIONS.get_interface(
-                self.config._experts_implementation, original_forward
-            )
+            experts_forward = experts_interface.get_interface(self.config._experts_implementation, original_forward)
             return experts_forward(self, *args, **kwargs)
 
         if not hasattr(experts_class, "_apply_gate"):

filename: src/transformers/quantizers/quantizer_finegrained_fp8.py
patch: @@ -77,10 +77,10 @@ def validate_environment(self, *args, **kwargs):
                 )
 
     def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:
-        from ..integrations.finegrained_fp8 import FP8Expert, FP8Linear
+        from ..integrations.finegrained_fp8 import FP8Experts, FP8Linear
 
         module, tensor_name = get_module_from_name(model, param_name)
-        if isinstance(module, (FP8Linear, FP8Expert)):
+        if isinstance(module, (FP8Linear, FP8Experts)):
             if self.pre_quantized or tensor_name == "bias":
                 return False
             else:
@@ -142,6 +142,10 @@ def is_serializable(self):
     def is_trainable(self) -> bool:
         return False
 
+    @property
+    def is_compileable(self) -> bool:
+        return True
+
     def get_quantize_ops(self):
         from ..integrations.finegrained_fp8 import Fp8Quantize
 

filename: src/transformers/utils/generic.py
patch: @@ -939,7 +939,7 @@ def wrapper(self, *args, **kwargs):
 class GeneralInterface(MutableMapping):
     """
     Dict-like object keeping track of a class-wide mapping, as well as a local one. Allows to have library-wide
-    modifications though the class mapping, as well as local modifications in a single file with the local mapping.
+    modifications through the class mapping, as well as local modifications in a single file with the local mapping.
     """
 
     # Class instance object, so that a call to `register` can be reflected into all other files correctly, even if

filename: tests/quantization/finegrained_fp8/test_fp8.py
patch: @@ -386,15 +386,18 @@ def test_compute_module_sizes(self):
         # we should at least have 1.5 times memory reduction in total
         assert model_size[""] > quantized_model_size[""] * 1.5
 
-    def test_quantized_moe_forward(self):
+    @parameterized.expand(["eager", "batched_mm", "grouped_mm"])
+    def test_quantized_moe_forward(self, experts_implementation):
         """
         Checks implicitly if the moe implementation is correct, i.e. it does not crash for cases
         where the indices go over `top_k` as shown within the Minimax M2 model
         """
         model = AutoModelForCausalLM.from_pretrained(
             "hf-internal-testing/MiniMax-M2-Tiny-FP8",  # single layer version
+            experts_implementation=experts_implementation,
             device_map=self.device_map,
         )
+        assert model.config._experts_implementation == experts_implementation
 
         tokenizer = AutoTokenizer.from_pretrained("MiniMaxAI/MiniMax-M2")
         messages = [