Skip to content

Add batched GPU NMF support with sklearn-compatible multiplicative-update (MU) and Fast-HALS coordinate-descent (CD) solvers - #9

Merged
OlivierBakker merged 1 commit into
TrynkaLab:mainfrom
JerryIshihara:features/nmf-gpu-kernal-cd-solver
Aug 5, 2026
Merged

Add batched GPU NMF support with sklearn-compatible multiplicative-update (MU) and Fast-HALS coordinate-descent (CD) solvers#9
OlivierBakker merged 1 commit into
TrynkaLab:mainfrom
JerryIshihara:features/nmf-gpu-kernal-cd-solver

Conversation

@JerryIshihara

Copy link
Copy Markdown

Changes

  • Add the cnmf.gpunmf package, split into:
    • shared GPU/runtime utilities
    • batched MU solver
    • batched Fast-HALS CD solver
    • fused Triton CUDA CD kernel
  • Add --solver mu|cd selection.
  • Add GPU engine configuration for device, dtype, TF32, compilation, and replicate batch size.
  • Support batched factorization while preserving one spectra output per cNMF replicate.
  • Enforce solver/loss compatibility: CD supports Frobenius loss only.
  • Preserve the existing CPU/sklearn path as the default.
  • Replace the previous nmf_gpu.py with focused modules.

@OlivierBakker
OlivierBakker merged commit 38389e5 into TrynkaLab:main Aug 5, 2026
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants