From 59b496ae71b91b0a70e0b0140ca7ce03ace127b1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aleksandar=20Samard=C5=BEi=C4=87?=
 <asamardzic@quansight.com>
Date: Wed, 26 Mar 2025 15:07:14 +0100
Subject: [PATCH] Update

[ghstack-poisoned]
---
 .../ATen/native/sparse/cuda/SparseSemiStructuredLinear.cu | 8 +-------
 .../ATen/native/sparse/cuda/SparseSemiStructuredOps.cu    | 8 +-------
 2 files changed, 2 insertions(+), 14 deletions(-)

diff --git a/aten/src/ATen/native/sparse/cuda/SparseSemiStructuredLinear.cu b/aten/src/ATen/native/sparse/cuda/SparseSemiStructuredLinear.cu
index 287938962a34..e74d71fe1aff 100644
--- a/aten/src/ATen/native/sparse/cuda/SparseSemiStructuredLinear.cu
+++ b/aten/src/ATen/native/sparse/cuda/SparseSemiStructuredLinear.cu
@@ -75,12 +75,6 @@ Tensor two_four_sgemm(
     using LayoutC = LayoutOutput;
     constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
 
-    using BiasTileThreadMap = cutlass::epilogue::threadblock::OutputTileThreadLayout<
-        ThreadblockShape,
-        WarpShape,
-        ElementC,
-        AlignmentC,
-        NumEVTEpilogueStages>;
     using OutputTileThreadMap = cutlass::epilogue::threadblock::OutputTileThreadLayout<
         ThreadblockShape,
         WarpShape,
@@ -94,7 +88,7 @@ Tensor two_four_sgemm(
         cutlass::epilogue::threadblock::VisitorScalarBroadcast<ElementC>;
     using BiasTensor =
         cutlass::epilogue::threadblock::VisitorColBroadcast<
-            BiasTileThreadMap,
+            OutputTileThreadMap,
             ElementC,
             cute::Stride<cute::_1, cute::_0, int64_t>>;
     using Bias = std::conditional_t<use_bias, BiasTensor, BiasScalar>;
diff --git a/aten/src/ATen/native/sparse/cuda/SparseSemiStructuredOps.cu b/aten/src/ATen/native/sparse/cuda/SparseSemiStructuredOps.cu
index 3b0f3a3170ca..35d6559b62ce 100644
--- a/aten/src/ATen/native/sparse/cuda/SparseSemiStructuredOps.cu
+++ b/aten/src/ATen/native/sparse/cuda/SparseSemiStructuredOps.cu
@@ -70,12 +70,6 @@ void spgemm_cutlass(
     using LayoutC = LayoutOutput;
     constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
 
-    using TensorCTileThreadMap = cutlass::epilogue::threadblock::OutputTileThreadLayout<
-        ThreadblockShape,
-        WarpShape,
-        ElementC,
-        AlignmentC,
-        NumEVTEpilogueStages>;
     using OutputTileThreadMap = cutlass::epilogue::threadblock::OutputTileThreadLayout<
         ThreadblockShape,
         WarpShape,
@@ -105,7 +99,7 @@ void spgemm_cutlass(
         cutlass::epilogue::threadblock::VisitorScalarBroadcast<ElementC>;
     using TensorCTensor =
         cutlass::epilogue::threadblock::VisitorColBroadcast<
-            TensorCTileThreadMap,
+            OutputTileThreadMap,
             ElementC,
             cute::Stride<cute::_1, cute::_0, int64_t>>;
     using TensorC = std::conditional_t<use_tensor_c, TensorCTensor, TensorCScalar>;