[PR #442] Add torch dispatcher op for tf32_hc_prenorm_gemm #450

New issue

Open

opened 2026-10-26 21:18:00 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 21:18:00 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/442
Author: @leevan
Created: 9/14/2026
Status: 🔄 Open

Base: main ← Head: feat/torchop-tf32-hc-prenorm

📝 Commits (2)

38f0dc2 Add torch dispatcher op for tf32_hc_prenorm_gemm
9f07225 Address review: FRAGMENT registration, num_splits check, test

📊 Changes

3 files changed (+96 additions, -0 deletions)

View changed files

➕ csrc/apis/torch_ops.hpp (+60 -0)
📝 csrc/python_api.cpp (+1 -0)
📝 tests/test_hyperconnection.py (+35 -0)

📄 Description

What

Registers deep_gemm::tf32_hc_prenorm_gemm as a torch dispatcher op (torch.ops.deep_gemm.tf32_hc_prenorm_gemm) in addition to the existing pybind binding.

Why

External C++ (and Python) callers can then invoke this kernel through the PyTorch dispatcher without including DeepGEMM headers or linking against DeepGEMM — e.g. a host-side C++ extension can resolve it via c10::Dispatcher::singleton().findSchemaOrThrow("deep_gemm::tf32_hc_prenorm_gemm", ""). The schema is compiled into the _C extension alongside the pybind bindings, so the op is available as soon as import deep_gemm loads _C.

Details

New header csrc/apis/torch_ops.hpp:
- A thin forwarder to deep_gemm::hyperconnection::tf32_hc_prenorm_gemm, converting std::optional<int64_t> → std::optional<int>.
- TORCH_LIBRARY(deep_gemm) schema: tf32_hc_prenorm_gemm(Tensor a, Tensor b, Tensor(d!) d, Tensor(s!) sqr_sum, int? num_splits=None) -> () — d and sqr_sum are marked mutable (written in place), num_splits defaults to None to match the pybind signature.
- TORCH_LIBRARY_IMPL(deep_gemm, CUDA) registering the CUDA implementation.
csrc/python_api.cpp: include the new header so the schema is compiled into _C.

Notes

The TORCH_LIBRARY handle is named lib (not m) so the .pyi stub generator, which scans csrc for pybind-style m.def(...), skips these dispatcher schemas.
Only a CUDA impl is registered (no Meta/fake kernel); the op is intended for eager execution. Happy to add a Meta kernel if there is interest in tracing through it under torch.compile.
This wraps a single API to establish the pattern; more host APIs can be exposed the same way later if useful.

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepGEMM/pull/442 **Author:** [@leevan](https://github.com/leevan) **Created:** 9/14/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `feat/torchop-tf32-hc-prenorm` --- ### 📝 Commits (2) - [`38f0dc2`](https://github.com/deepseek-ai/DeepGEMM/commit/38f0dc20f933398ab52393bd5c57dd35011c5c84) Add torch dispatcher op for tf32_hc_prenorm_gemm - [`9f07225`](https://github.com/deepseek-ai/DeepGEMM/commit/9f0722503166c4edb5f62e2841b24a1c957ca8ba) Address review: FRAGMENT registration, num_splits check, test ### 📊 Changes **3 files changed** (+96 additions, -0 deletions) <details> <summary>View changed files</summary> ➕ `csrc/apis/torch_ops.hpp` (+60 -0) 📝 `csrc/python_api.cpp` (+1 -0) 📝 `tests/test_hyperconnection.py` (+35 -0) </details> ### 📄 Description ## What Registers `deep_gemm::tf32_hc_prenorm_gemm` as a **torch dispatcher op** (`torch.ops.deep_gemm.tf32_hc_prenorm_gemm`) in addition to the existing pybind binding. ## Why External C++ (and Python) callers can then invoke this kernel through the PyTorch dispatcher **without including DeepGEMM headers or linking against DeepGEMM** — e.g. a host-side C++ extension can resolve it via `c10::Dispatcher::singleton().findSchemaOrThrow("deep_gemm::tf32_hc_prenorm_gemm", "")`. The schema is compiled into the `_C` extension alongside the pybind bindings, so the op is available as soon as `import deep_gemm` loads `_C`. ## Details - New header `csrc/apis/torch_ops.hpp`: - A thin forwarder to `deep_gemm::hyperconnection::tf32_hc_prenorm_gemm`, converting `std::optional<int64_t>` → `std::optional<int>`. - `TORCH_LIBRARY(deep_gemm)` schema: `tf32_hc_prenorm_gemm(Tensor a, Tensor b, Tensor(d!) d, Tensor(s!) sqr_sum, int? num_splits=None) -> ()` — `d` and `sqr_sum` are marked mutable (written in place), `num_splits` defaults to `None` to match the pybind signature. - `TORCH_LIBRARY_IMPL(deep_gemm, CUDA)` registering the CUDA implementation. - `csrc/python_api.cpp`: include the new header so the schema is compiled into `_C`. ### Notes - The `TORCH_LIBRARY` handle is named `lib` (not `m`) so the `.pyi` stub generator, which scans csrc for pybind-style `m.def(...)`, skips these dispatcher schemas. - Only a CUDA impl is registered (no Meta/fake kernel); the op is intended for eager execution. Happy to add a Meta kernel if there is interest in tracing through it under `torch.compile`. - This wraps a single API to establish the pattern; more host APIs can be exposed the same way later if useful. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-26 21:18:00 +00:00

bsakaguchi referenced this issue

2026-10-26 21:18:02 +00:00

[PR #450] feat(gemm): opt-in labels-contract checker for m-grouped contiguous GEMM #454

bsakaguchi referenced this issue

2026-10-26 21:18:03 +00:00

[PR #447] [MERGED] Integrate main APIs and SM120 kernels while preserving nv_dev features #455