migrated DeepGEMM to torch stable ABI (unchecked) by cleonard530 · Pull Request #8 · cleonard530/DeepGEMM

cleonard530 · 2026-09-18T20:19:22Z

migration to stable abi (new with deep_jit)

cleonard530 · 2026-09-21T15:01:13Z

+// Build integer indices entirely on-device so layout conversion remains graph-capturable.
+inline torch::stable::Tensor arange(int64_t end, torch::stable::Device device) {
+    auto ones = torch::stable::full({end}, 1, torch::headeronly::ScalarType::Long, std::nullopt, device);
+    std::array<StableIValue, 3> stack{
+        torch::stable::detail::from(ones), torch::stable::detail::from(int64_t{0}),
+        torch::stable::detail::from(std::optional<torch::headeronly::ScalarType>{})};
+    detail::call_dispatcher("aten::cumsum", "", stack.data());
+    const auto indices = torch::stable::detail::to<torch::stable::Tensor>(stack[0]);
+    return torch::stable::subtract(indices, ones);
+}


This is inefficient, we will need to find a a better way of doing this

9a86ae2 re-ported the SM90 paged MQA kv_block=32/next_n=4 support from deepseek-ai/DeepGEMM's nv_dev branch onto the 26/09 release base. nv_dev carried four guards for zero context lengths that the 26/09 base never had, and the re-port dropped all of them. With every context length zero, `sum` is zero, so the binary search in `sm90_paged_mqa_logits_metadata` finds no prefix sum greater than `seg_starts` and returns `lo == batch_size`. Using that unclamped as `q_idx` reads `prefix_sum[batch_size]`, which is one past the end of a shared buffer sized to exactly `align(batch_size, 32)` ints. When `batch_size` is already 32-aligned the read is out of range and the kernel traps: CUDBG_EXCEPTION_WARP_OUT_OF_RANGE_ADDRESS (5) #0 deep_gemm::sched::sm90_paged_mqa_logits_metadata<1024u, 256u, 132u, false> vLLM hits this on every DeepSeek V4 startup: the DSA indexer warmup runs a dummy batch of max_num_seqs (default 1024) rows with all context lengths zero. Restore the four guards: - `total == 0` early-returns with one-past-the-end sentinels, avoiding the `prefix_sum[batch_size]` read and keeping a trailing SM off a synthetic zero-KV task. - `q_idx` is clamped to `batch_size - 1` before indexing prefix_sum. - The scheduler constructor checks `exist_q_atom_idx` before `refresh_num_kv_and_advance` dereferences context_lens/indices, and initializes current_advance/current_num_kv/last_advance for the empty range. - `fetch_next_task` skips a zero-context request's atoms rather than exposing an empty task, for the case where traversal crosses one between two non-empty requests. The scheduler is now behaviorally identical to the nv_dev version it was ported from. tests/test_attention.py: `test_paged_mqa_logits` draws context lengths around a positive average and so never produces an empty request, which is why this regressed silently. Add `test_paged_mqa_logits_zero_context` covering all-zero and interleaved-zero context lengths over block_kv 32/64, next_n 1/2/4, and batch sizes 32 and 1024 (both 32-aligned, so `prefix_sum[batch_size]` is exactly one past the end). Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

cleonard530 · 2026-09-29T13:41:23Z

-TORCH_LIBRARY(deep_gemm, m) {
-    // Register JIT objects
-    m.class_<deep_gemm::JitRuntimeHandle>("Runtime");
-    m.def("get_jit() -> __torch__.torch.classes.deep_gemm.Runtime", []() { return c10::make_intrusive<deep_gemm::JitRuntimeHandle>(); });
+STABLE_TORCH_LIBRARY(deep_gemm, m) {
+    m.def("_ensure_jit() -> ()");
+}
+STABLE_TORCH_LIBRARY_IMPL(deep_gemm, CompositeExplicitAutograd, m) {
+    m.impl("_ensure_jit", TORCH_BOX(&deep_gemm::ensure_jit));


JitRuntimeHandle was removed because the PyTorch stable ABI does not support registering custom C++ classes. The Python Runtime wrapper preserves the existing get_jit() API while _ensure_jit() still initializes DeepJIT; no runtime methods are lost because DeepJIT’s Python binding did not expose any.

cleonard530 · 2026-09-29T13:45:25Z

+class Runtime:
+    """Opaque compatibility type formerly returned by the pybind custom class."""
+
+
+def get_jit() -> Runtime:
+    """Initialize and return the opaque runtime handle kept by the legacy API."""
+    _torch_ops._ensure_jit()
+    return Runtime()


The wrapper to preserve the existing get_jit() API.

The Python Runtime placeholder preserves the legacy API’s surface semantics, but not the original Torch custom-class type identity. It is intentionally empty because DeepJIT exposed no runtime methods through Python.

cleonard530 · 2026-09-30T13:49:31Z

+    values = _torch_ops._get_bf16_mega_gate_config(num_tokens, hidden, num_routed_experts, num_topk)
+    return dict(zip(('block_tokens', 'num_mma_ctas', 'num_split_k',
+                     'num_expert_groups', 'num_gate_warpgroups', 'num_sms'), values))


STABLE_TORCH_LIBRARY cannot register a dict, so a list was registered instead and this wraps the dict around the list values.

cleonard530 · 2026-09-30T14:06:33Z


    // Create the cuBLASLt handle ourselves
    cublasLtHandle_t cublaslt_handle;
-    bool use_pytorch_managed_cublaslt_handle;


PyTorch does not have a stable ABI equivalent of at::cuda::getCurrentCUDABlasLtHandle() so we remove this option. I have not been able to find any test that use this and by default it is not used, so I don't know of any major impact, but this should be looked at closely.

cleonard530 · 2026-09-30T14:43:27Z

-    static std::unordered_map<c10::cuda::CUDAStream, torch::Tensor> workspaces;
+    const auto stream = torch_compat::stream_key(reference);
+
+    static std::map<torch_compat::StreamKey, torch::stable::Tensor> workspaces;


std::pair<int32_t, cudaStream_t> (StreamKey) does not have hashing support.

cleonard530 · 2026-09-30T14:46:43Z

        DG_HOST_ASSERT(not disable_ue8m0_cast);
        const auto broadcasted = gran_mn == 1 ? sf :
-                                 sf.index_select(-2, torch::arange(mn, at::TensorOptions().device(sf.device())).floor_divide_(gran_mn));
+                                 torch_compat::index_select(sf, -2, torch_compat::floor_divide(deep_gemm_arange_int64(mn, sf.device()), gran_mn));


There is no stable shim for torch arange, so we implemented a deep_gemm kernel for this. This now uses DeepJIT for compilation on the first use.

cleonard530 · 2026-09-30T14:56:03Z

-static c10::Dict<std::string, int64_t> get_bf16_mega_gate_config(const int& num_tokens, const int& hidden,
+static std::vector<int64_t> get_bf16_mega_gate_config(const int& num_tokens, const int& hidden,


c10::Dict<std::string, int64_t> cannot be used in a stable operator schema. The stable ABI supports int[], so the C++ implementation returns std::vector<int64_t>. The Python wrapper converts that vector back into the public dictionary

cleonard530 · 2026-09-30T15:05:52Z

-TORCH_LIBRARY_IMPL(deep_gemm, CatchAll, m) {
-    m.impl("get_bf16_mega_gate_config", TORCH_FN(deep_gemm::torch_registration::get_bf16_mega_gate_config));
+STABLE_TORCH_LIBRARY_IMPL(deep_gemm, CompositeExplicitAutograd, m) {
+    m.impl("_get_bf16_mega_gate_config", TORCH_BOX(&deep_gemm::torch_registration::get_bf16_mega_gate_config));
 }


CompositeExplicitAutograd behaves the same as Catchall for get_bf16_mega_gate_config, which has no tensor arguments.

Also, Added underscore (_) to torch ops registration name to indicate difference in return type between it and python wrapper (List vs. dict)

cleonard530 · 2026-09-30T16:01:25Z

-        DG_HOST_ASSERT(c10::cuda::currentStreamCaptureStatusMayInitCtx() == c10::cuda::CaptureStatus::None);
-        workspace = torch::zeros({kNumWorkspaceBytes}, options.dtype(torch::kByte));
+        DG_HOST_ASSERT(not torch_compat::is_capturing(stream.second));
+        workspace = torch::stable::new_zeros(reference, {kNumWorkspaceBytes}, torch::headeronly::ScalarType::Byte);
    }


The new not torch_compat::is_capturing(stream.second) check is more explicit and stable-ABI compatible: it checks capture status on the exact CUDA stream associated with the reference tensor’s device. The previous c10::cuda::currentStreamCaptureStatusMayInitCtx() == c10::cuda::CaptureStatus::None call used an unstable PyTorch API and queried PyTorch’s current stream instead.

Co-authored-by: codex Signed-off-by: Chris Leonard <chleonar@redhat.com>

…table arange method. Fixed merge issues. removed unnecessaary torch_compat::from_blob and megred 3 namespace blocks into one (they were probbaly disjoint before due to Agent iteration) Co-authored-by: codex Signed-off-by: Chris Leonard <chleonar@redhat.com>

…ion. Co-authored-by: codex Signed-off-by: Chris Leonard <chleonar@redhat.com>

…M100, and SM120 heuristics to avoid repeated stable ABI lookups. Co-authored-by: codex Signed-off-by: Chris Leonard <chleonar@redhat.com>

cleonard530 commented Sep 21, 2026

View reviewed changes

cleonard530 force-pushed the migrate_pybind_to_torch_library_dev branch from b201a73 to dec2595 Compare September 24, 2026 12:56

cleonard530 force-pushed the migrate_pybind_to_torch_library_dev branch 2 times, most recently from 27cc188 to 496fd0b Compare September 28, 2026 23:15

cleonard530 force-pushed the migrate_deep_gemm_to_stable_abi_2 branch 2 times, most recently from 99b10e5 to 5fefca5 Compare September 29, 2026 12:46

cleonard530 commented Sep 29, 2026

View reviewed changes

cleonard530 force-pushed the migrate_pybind_to_torch_library_dev branch from 9560d13 to 45d9a65 Compare September 30, 2026 12:27

cleonard530 force-pushed the migrate_deep_gemm_to_stable_abi_2 branch from 5fefca5 to f07c95c Compare September 30, 2026 13:23

cleonard530 commented Sep 30, 2026

View reviewed changes

cleonard530 force-pushed the migrate_deep_gemm_to_stable_abi_2 branch from f07c95c to 57a5da2 Compare September 30, 2026 17:13

cleonard530 added 4 commits September 30, 2026 17:15

migrated DeepGEMM to torch stable ABI (unchecked)

286f765

Co-authored-by: codex Signed-off-by: Chris Leonard <chleonar@redhat.com>

moved min torch version from 2.13 to 2.10

b06ddc9

Co-authored-by: codex Signed-off-by: Chris Leonard <chleonar@redhat.com>

Clean up stable ABI compatibility wrappers and strided tensor allocat…

4bba562

…ion. Co-authored-by: codex Signed-off-by: Chris Leonard <chleonar@redhat.com>

cleonard530 force-pushed the migrate_deep_gemm_to_stable_abi_2 branch from 57a5da2 to 4bba562 Compare September 30, 2026 17:17

cleonard530 mentioned this pull request Sep 30, 2026

Migrate deep gemm to stable abi vllm-project/DeepGEMM#21

Merged

Cache GEMM element sizes in descriptors and reuse them across SM90, S…

059e168

…M100, and SM120 heuristics to avoid repeated stable ABI lookups. Co-authored-by: codex Signed-off-by: Chris Leonard <chleonar@redhat.com>

Navigation Menu

Repository navigation

migrated DeepGEMM to torch stable ABI (unchecked) - #8

migrated DeepGEMM to torch stable ABI (unchecked)#8
cleonard530 wants to merge 5 commits into
migrate_pybind_to_torch_library_devfrom
migrate_deep_gemm_to_stable_abi_2

cleonard530 commented Sep 18, 2026

Uh oh!

cleonard530 Sep 21, 2026

Uh oh!

cleonard530 Sep 29, 2026

Uh oh!

cleonard530 Sep 29, 2026

Uh oh!

cleonard530 Sep 30, 2026

Uh oh!

cleonard530 Sep 30, 2026

Uh oh!

cleonard530 Sep 30, 2026

Uh oh!

cleonard530 Sep 30, 2026

Uh oh!

cleonard530 Sep 30, 2026

Uh oh!

cleonard530 Sep 30, 2026 •

edited

Loading

Uh oh!

cleonard530 Sep 30, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

		static c10::Dict<std::string, int64_t> get_bf16_mega_gate_config(const int& num_tokens, const int& hidden,
		static std::vector<int64_t> get_bf16_mega_gate_config(const int& num_tokens, const int& hidden,

Conversation

cleonard530 commented Sep 18, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cleonard530 Sep 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

cleonard530 Sep 30, 2026 •

edited

Loading