Fix deepseek-v3 with torch.compile in PyTorch 2.6. (#5213)

76f44c2a · Richard Zou · GitHub · 1078396f · 76f44c2a
Unverified Commit 76f44c2a authored Apr 10, 2025 by Richard Zou Committed by GitHub Apr 10, 2025
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

sgl-kernel/csrc/common_extension.cc sgl-kernel/csrc/common_extension.cc +2 -1

No files found.
--- a/sgl-kernel/csrc/common_extension.cc
+++ b/sgl-kernel/csrc/common_extension.cc
@@ -177,7 +177,8 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
   */
  m.def(
      "bmm_fp8(Tensor A, Tensor B, Tensor! D, Tensor A_scale, Tensor B_scale, Tensor workspace_buffer, int "
-      "cublas_handle, int cuda_stream) -> ()");
+      "cublas_handle, int cuda_stream) -> ()",
+      {at::Tag::needs_fixed_stride_order});
  m.impl("bmm_fp8", torch::kCUDA, &bmm_fp8);

  m.def(