NVIDIA / NVIDIA/TransformerEngine

[JAX] Remove D2H in GroupedQuantizeFFI/GroupedGemmFFI

Open
#2,469 1 comment 0 reactions 1 assignee View on GitHub

@jberchtold-nvidia is already working on this.

Since Jan 23, 2026.

2.17 MoE
Dominant language
Python
Stars
3.5k
Forks
831
Avg merge
3d 11h
Merged PRs (30d)
65

Description

Eliminate device-to-host (D2H) copies in JAX GroupedQuantizeFFI and GroupedGemmFFI APIs, providing a fully device-resident workflow for MoE quantization/training.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.