[JAX] grouped_gemm() uses variadic arguments - #1658

Merged

phu0ngng merged 14 commits into

NVIDIA:mainfrom

huanghua1994:JAX-GroupedGEMM-VariadicArgs

Apr 14, 2025

Commits on Apr 11, 2025

grouped_add example to test variadic args
huanghua1994
committed
2f7342a
New GroupedGemmPrimitive using variadic args
huanghua1994
committed
59b17fd
Remove NVTE_NO_SCALING in new grouped_gemm
huanghua1994
committed
937f98d
Remove squeeze() to reduce D2D memcpy
huanghua1994
committed
963fd4c
Clean up code
huanghua1994
committed
408acde
.
huanghua1994
committed
167ef32
[pre-commit.ci] auto fixes from pre-commit.com hooks

pre-commit-ci[bot]
authored and
huanghua1994
committed
eceff90
Remove unused file
huanghua1994
committed
c1e33bd
Revert to the list append fashion to simplify code
huanghua1994
committed
538451d
[pre-commit.ci] auto fixes from pre-commit.com hooks

pre-commit-ci[bot]
authored and
huanghua1994
committed
fe7c3bf
Rebase and update; fix scale_inv shapes in C++
huanghua1994
committed
e5bee41
[pre-commit.ci] auto fixes from pre-commit.com hooks

pre-commit-ci[bot]
authored and
huanghua1994
committed
23ae948
Bug fix and minor revision
huanghua1994
committed
c850c00

Commits on Apr 14, 2025

Merge branch 'main' into JAX-GroupedGEMM-VariadicArgs
phu0ngng
authored
9a17a8c