Skip to content

Make paged MQA metadata independent of PyTorch C++ ABI - #70

Draft
nvpohanh wants to merge 1 commit into
sgl-project:devfrom
nvpohanh:codex/nvbug-6317902-tvm-ffi
Draft

Make paged MQA metadata independent of PyTorch C++ ABI#70
nvpohanh wants to merge 1 commit into
sgl-project:devfrom
nvpohanh:codex/nvbug-6317902-tvm-ffi

Conversation

@nvpohanh

@nvpohanh nvpohanh commented Aug 3, 2026

Copy link
Copy Markdown

[by Codex]

Why this is needed

sgl-deep-gemm distributes a prebuilt _C.so, but the paged-MQA metadata API currently converts TVM-FFI tensors to torch::Tensor with torch::from_blob. This unnecessarily couples that path to the PyTorch C++ ABI used when the wheel was built. PyTorch does not guarantee that extension binaries using ATen object internals remain compatible across releases.

The metadata launcher only needs tensor metadata, raw device pointers, and an output allocation. Keeping this path on the TVM-FFI boundary avoids constructing ATen objects and makes the boundary match the data actually consumed by the CUDA launcher.

What this PR does

This PR removes PyTorch C++ tensor objects from the exported get_paged_mqa_logits_metadata path:

  • validates shape, dtype, contiguity, and device with tvm::ffi::TensorView;
  • allocates the output with Tensor::FromEnvAlloc(TVMFFIEnvTensorAlloc, ...);
  • passes raw CUDA pointers and scalar metadata to the existing SM90, SM100, and SM120 launchers;
  • returns the TVM-FFI tensor directly instead of converting an ATen tensor through DLPack;
  • keeps compatibility wrappers for existing internal ATen call sites; and
  • adds direct metadata tests covering several 2-D input shapes.

This is intentionally an incremental patch. Other exported DeepGEMM paths still use ATen and the wheel still links libtorch. Those paths should be migrated to the same TVM-FFI/raw-pointer pattern before claiming that the entire wheel is independent of the PyTorch C++ ABI.

Validation

  • compiled and exercised the updated SM120 metadata path;
  • added direct coverage for four 2-D context_lens shapes;
  • checked output shape, dtype, device, and contiguity; and
  • verified Python syntax and git diff --check.

The complete DeepGEMM test suite has not been run across all supported GPU architectures yet, so this is opened as a draft.

@nvpohanh

nvpohanh commented Aug 3, 2026

Copy link
Copy Markdown
Author

@Fridge003 Could you check if this is a reasonable fix?

Context: we are upgrading the PyTorch version in our NGC SGLang container to prepare for Rubin support, but we ran into a PyTorch ABI issue caused by sgl-deep-geem. This is the agent's suggested fix. I will take a closer look tomorrow.

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.

1 participant