Skip to content

build: lower CUDA floor to 12.9, stage oversized GEMM smem dynamically #235

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
teo-mateo wants to merge 1 commit into
base: master
Choose a base branch
from
Draft

build: lower CUDA floor to 12.9, stage oversized GEMM smem dynamically #235

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,9 +33,9 @@ if(CMAKE_GENERATOR MATCHES "Ninja")
set(CMAKE_JOB_POOL_LINK cuda_link)
endif()

if(CMAKE_CUDA_COMPILER_VERSION VERSION_LESS 13.1)
if(CMAKE_CUDA_COMPILER_VERSION VERSION_LESS 12.9)
message(FATAL_ERROR
"NInfer requires CUDA 13.1 or newer; found CUDA compiler "
"NInfer requires CUDA 12.9 or newer; found CUDA compiler "
"${CMAKE_CUDA_COMPILER_VERSION}")
endif()

Expand Down
17 changes: 13 additions & 4 deletions src/ops/attn_input_proj/nvfp4/nvfp4_attn_input_w4a4.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,10 +71,19 @@ void launch_gemm(const Weight& weight, Tensor& q, Tensor& gate, Tensor& k, Tenso
static_cast<__nv_bfloat16*>(v.data),
};
const float alpha = 1.0F / (weight.input_scale_divisor * weight.weight_scale_divisor);
nvfp4_w4a4_mma_kernel<Geometry, Schedule><<<grid, Schedule::kThreads, 0, stream>>>(
activation, static_cast<const std::uint8_t*>(weight.qdata),
static_cast<const std::uint8_t*>(weight.scales), tokens, alpha, Nvfp4IdentityEpilogue{},
output);
constexpr std::size_t kDynamicBytes = nvfp4_w4a4_mma_dynamic_bytes<Schedule>();
if constexpr (kDynamicBytes > kNvfp4W4a4StaticSharedBytes) {
static const cudaError_t attribute = cudaFuncSetAttribute(
nvfp4_w4a4_mma_kernel<Geometry, Schedule, Nvfp4IdentityEpilogue,
Nvfp4W4a4AttentionOutput>,
cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(kDynamicBytes));
CUDA_CHECK(attribute);
}
nvfp4_w4a4_mma_kernel<Geometry, Schedule>
<<<grid, Schedule::kThreads, kDynamicBytes, stream>>>(
activation, static_cast<const std::uint8_t*>(weight.qdata),
static_cast<const std::uint8_t*>(weight.scales), tokens, alpha,
Nvfp4IdentityEpilogue{}, output);
CUDA_CHECK(cudaGetLastError());
}

Expand Down
16 changes: 12 additions & 4 deletions src/ops/attn_input_proj/w8/w8_attn_input_gemm_mma.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,11 +16,19 @@ using CompanionOutput = W8SplitOutput3<4096, 1024, 1024>;
template <class Schedule, bool Full, int Rows, class Output>
void launch_variant(const Tensor& x, const Weight& weight, Output output, cudaStream_t stream) {
const dim3 grid(Rows / Schedule::BM, static_cast<unsigned>(div_up(x.ne[1], Schedule::BN)), 1u);
constexpr std::size_t kDynamicBytes = w8_rowsplit_gemm_mma_dynamic_bytes<Schedule>();
if constexpr (kDynamicBytes > kW8SmallTMmaStaticSharedBytes) {
static const cudaError_t attribute = cudaFuncSetAttribute(
w8_rowsplit_gemm_mma_kernel<Schedule, Full, W8Epilogue::Store, Output>,
cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(kDynamicBytes));
CUDA_CHECK(attribute);
}
w8_rowsplit_gemm_mma_kernel<Schedule, Full, W8Epilogue::Store, Output>
<<<grid, Schedule::THREADS, 0, stream>>>(static_cast<const __nv_bfloat16*>(x.data),
static_cast<const std::uint8_t*>(weight.qdata),
static_cast<const std::uint8_t*>(weight.scales),
output, Rows, kHidden, x.ne[1], kHidden);
<<<grid, Schedule::THREADS, kDynamicBytes, stream>>>(
static_cast<const __nv_bfloat16*>(x.data),
static_cast<const std::uint8_t*>(weight.qdata),
static_cast<const std::uint8_t*>(weight.scales), output, Rows, kHidden, x.ne[1],
kHidden);
}

template <class Schedule, int Rows, class Output>
Expand Down
33 changes: 30 additions & 3 deletions src/ops/attn_input_proj/w8/w8_attn_input_gemm_splitk.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,8 +36,17 @@ void launch_output(const Tensor& x, const Weight& weight, Output output, cudaStr
: 48;
using Geometry = W8LinearGeometry<Rows, kHidden>;
using Schedule = W8SmallTMmaDefaultSchedule<TileCols, ActiveCols>;
constexpr std::size_t kDynamicBytes =
w8_small_t_mma_dynamic_bytes<Schedule, false, ActiveCols>();
if constexpr (kDynamicBytes > kW8SmallTMmaStaticSharedBytes) {
static const cudaError_t attribute = cudaFuncSetAttribute(
w8_small_t_mma_kernel<Geometry, ActiveCols, Schedule, Output,
W8SmallTMmaStoreEpilogue, W8SmallTMmaIdentityRows>,
cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(kDynamicBytes));
CUDA_CHECK(attribute);
}
w8_small_t_mma_kernel<Geometry, ActiveCols, Schedule>
<<<Rows / kRowsPerCta, Schedule::kThreads, 0, stream>>>(
<<<Rows / kRowsPerCta, Schedule::kThreads, kDynamicBytes, stream>>>(
static_cast<const __nv_bfloat16*>(x.data),
static_cast<const std::uint8_t*>(weight.qdata),
static_cast<const std::uint8_t*>(weight.scales), output);
Expand Down Expand Up @@ -87,8 +96,17 @@ void launch_target_medium_cols(const Tensor& x, const Weight& weight, Tensor& q,
const TargetOutput output{
static_cast<__nv_bfloat16*>(q.data), static_cast<__nv_bfloat16*>(k.data),
static_cast<__nv_bfloat16*>(gate.data), static_cast<__nv_bfloat16*>(v.data)};
constexpr std::size_t kDynamicBytes =
w8_rowsplit_medium_t_splitk_dynamic_bytes<KSplits, TileCols, NGroups>();
if constexpr (kDynamicBytes > kW8SmallTMmaStaticSharedBytes) {
static const cudaError_t attribute = cudaFuncSetAttribute(
w8_rowsplit_medium_t_splitk_kernel<kHidden, TileCols, KSplits, NGroups, MinBlocks,
TargetOutput>,
cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(kDynamicBytes));
CUDA_CHECK(attribute);
}
w8_rowsplit_medium_t_splitk_kernel<kHidden, TileCols, KSplits, NGroups, MinBlocks>
<<<kTargetRows / kRowsPerCta, KSplits * NGroups * 32, 0, stream>>>(
<<<kTargetRows / kRowsPerCta, KSplits * NGroups * 32, kDynamicBytes, stream>>>(
static_cast<const __nv_bfloat16*>(x.data),
static_cast<const std::uint8_t*>(weight.qdata),
static_cast<const std::uint8_t*>(weight.scales), output, x.ne[1]);
Expand All @@ -101,8 +119,17 @@ void launch_companion_medium_cols(const Tensor& x, const Weight& weight, Tensor&
const CompanionOutput output{static_cast<__nv_bfloat16*>(q.data),
static_cast<__nv_bfloat16*>(k.data),
static_cast<__nv_bfloat16*>(v.data)};
constexpr std::size_t kDynamicBytes =
w8_rowsplit_medium_t_splitk_dynamic_bytes<KSplits, TileCols, NGroups>();
if constexpr (kDynamicBytes > kW8SmallTMmaStaticSharedBytes) {
static const cudaError_t attribute = cudaFuncSetAttribute(
w8_rowsplit_medium_t_splitk_kernel<kHidden, TileCols, KSplits, NGroups, MinBlocks,
CompanionOutput>,
cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(kDynamicBytes));
CUDA_CHECK(attribute);
}
w8_rowsplit_medium_t_splitk_kernel<kHidden, TileCols, KSplits, NGroups, MinBlocks>
<<<kCompanionRows / kRowsPerCta, KSplits * NGroups * 32, 0, stream>>>(
<<<kCompanionRows / kRowsPerCta, KSplits * NGroups * 32, kDynamicBytes, stream>>>(
static_cast<const __nv_bfloat16*>(x.data),
static_cast<const std::uint8_t*>(weight.qdata),
static_cast<const std::uint8_t*>(weight.scales), output, x.ne[1]);
Expand Down
20 changes: 18 additions & 2 deletions src/ops/attn_input_proj/w8/w8_dflash2_attn_input.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,9 +42,18 @@ void launch_small(const Tensor& x, const Weight& weight, Tensor& q, Tensor& k, T
const Output output{static_cast<__nv_bfloat16*>(q.data), static_cast<__nv_bfloat16*>(k.data),
static_cast<__nv_bfloat16*>(v.data)};
constexpr int kBlocks = Geometry::kOutputRows / Schedule::kRowsPerCta;
constexpr std::size_t kDynamicBytes =
w8_small_t_mma_dynamic_bytes<Schedule, !Exact, Columns>();
if constexpr (kDynamicBytes > kW8SmallTMmaStaticSharedBytes) {
static const cudaError_t attribute = cudaFuncSetAttribute(
w8_small_t_mma_kernel<Geometry, Columns, Schedule, Output, W8SmallTMmaStoreEpilogue,
W8SmallTMmaIdentityRows, false, !Exact>,
cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(kDynamicBytes));
CUDA_CHECK(attribute);
}
w8_small_t_mma_kernel<Geometry, Columns, Schedule, Output, W8SmallTMmaStoreEpilogue,
W8SmallTMmaIdentityRows, false, !Exact>
<<<kBlocks, Schedule::kThreads, 0, stream>>>(
<<<kBlocks, Schedule::kThreads, kDynamicBytes, stream>>>(
static_cast<const __nv_bfloat16*>(x.data),
static_cast<const std::uint8_t*>(weight.qdata),
static_cast<const std::uint8_t*>(weight.scales), output, W8SmallTMmaStoreEpilogue{},
Expand All @@ -69,8 +78,15 @@ void launch_mma_slice(const Tensor& x, const Weight& weight, Tensor& q, Tensor&
static_cast<__nv_bfloat16*>(v.data)};
const dim3 grid(Geometry::kOutputRows / Schedule::BM,
static_cast<unsigned>(div_up(x.ne[1], Schedule::BN)), 1u);
constexpr std::size_t kDynamicBytes = w8_rowsplit_gemm_mma_dynamic_bytes<Schedule>();
if constexpr (kDynamicBytes > kW8SmallTMmaStaticSharedBytes) {
static const cudaError_t attribute = cudaFuncSetAttribute(
w8_rowsplit_gemm_mma_kernel<Schedule, Full, W8Epilogue::Store, Output>,
cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(kDynamicBytes));
CUDA_CHECK(attribute);
}
w8_rowsplit_gemm_mma_kernel<Schedule, Full, W8Epilogue::Store, Output>
<<<grid, Schedule::THREADS, 0, stream>>>(
<<<grid, Schedule::THREADS, kDynamicBytes, stream>>>(
static_cast<const __nv_bfloat16*>(x.data),
static_cast<const std::uint8_t*>(weight.qdata),
static_cast<const std::uint8_t*>(weight.scales), output, Geometry::kOutputRows,
Expand Down
9 changes: 5 additions & 4 deletions src/ops/dynamic_grouped_conv/w8/w8_dynamic_grouped_conv_add_materialized.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,20 +40,21 @@ void tiled_projection(const Tensor& x, const Weight& weight, Tensor& out, cudaSt
using Geometry = W8LinearGeometry<kRows, InputRows>;
using Schedule = W8SmallTMmaSchedule<Warps, TileColumns, Warps == 8 ? 2 : 3,
W8SmallTMmaScaleAccess::Shared, Activation>;
constexpr int SharedBytes = TileColumns > 64 ? sizeof(W8SmallTMmaSharedStorage<Schedule>) : 0;
if constexpr (SharedBytes > 0) {
constexpr std::size_t kDynamicBytes =
w8_small_t_mma_dynamic_bytes<Schedule, true, TileColumns>();
if constexpr (kDynamicBytes > kW8SmallTMmaStaticSharedBytes) {
static const cudaError_t attribute = cudaFuncSetAttribute(
w8_small_t_mma_kernel<Geometry, TileColumns, Schedule, W8ContiguousOutput,
W8SmallTMmaStoreEpilogue, W8SmallTMmaIdentityRows, false, true>,
cudaFuncAttributeMaxDynamicSharedMemorySize, SharedBytes);
cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(kDynamicBytes));
CUDA_CHECK(attribute);
}
const int columns = x.ne[1];
W8ContiguousOutput output{static_cast<__nv_bfloat16*>(out.data), kRows};
const dim3 grid(kRows / 16, (columns + TileColumns - 1) / TileColumns);
w8_small_t_mma_kernel<Geometry, TileColumns, Schedule, W8ContiguousOutput,
W8SmallTMmaStoreEpilogue, W8SmallTMmaIdentityRows, false, true>
<<<grid, Schedule::kThreads, SharedBytes, stream>>>(
<<<grid, Schedule::kThreads, kDynamicBytes, stream>>>(
static_cast<const __nv_bfloat16*>(x.data),
static_cast<const std::uint8_t*>(weight.qdata),
static_cast<const std::uint8_t*>(weight.scales), output, W8SmallTMmaStoreEpilogue{},
Expand Down
19 changes: 14 additions & 5 deletions src/ops/gdn_input_proj/nvfp4/nvfp4_gdn_input_w4a4.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,11 +25,20 @@ void launch_gemm(const Weight& weight, Tensor& qkv, Tensor& z, Nvfp4W4a4Workspac
(tokens + Schedule::kBlockM - 1) / Schedule::kBlockM);
const Nvfp4W4a4MaterializedActivation activation{workspace.codes, workspace.scales};
const float alpha = 1.0F / (weight.input_scale_divisor * weight.weight_scale_divisor);
nvfp4_w4a4_mma_kernel<Geometry, Schedule><<<grid, Schedule::kThreads, 0, stream>>>(
activation, static_cast<const std::uint8_t*>(weight.qdata),
static_cast<const std::uint8_t*>(weight.scales), tokens, alpha, Nvfp4IdentityEpilogue{},
Nvfp4GdnInputOutput{static_cast<__nv_bfloat16*>(qkv.data),
static_cast<__nv_bfloat16*>(z.data)});
constexpr std::size_t kDynamicBytes = nvfp4_w4a4_mma_dynamic_bytes<Schedule>();
if constexpr (kDynamicBytes > kNvfp4W4a4StaticSharedBytes) {
static const cudaError_t attribute = cudaFuncSetAttribute(
nvfp4_w4a4_mma_kernel<Geometry, Schedule, Nvfp4IdentityEpilogue, Nvfp4GdnInputOutput>,
cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(kDynamicBytes));
CUDA_CHECK(attribute);
}
nvfp4_w4a4_mma_kernel<Geometry, Schedule>
<<<grid, Schedule::kThreads, kDynamicBytes, stream>>>(
activation, static_cast<const std::uint8_t*>(weight.qdata),
static_cast<const std::uint8_t*>(weight.scales), tokens, alpha,
Nvfp4IdentityEpilogue{},
Nvfp4GdnInputOutput{static_cast<__nv_bfloat16*>(qkv.data),
static_cast<__nv_bfloat16*>(z.data)});
CUDA_CHECK(cudaGetLastError());
}

Expand Down
16 changes: 12 additions & 4 deletions src/ops/gdn_input_proj/w8/w8_gdn_input_gemm_mma.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,11 +18,19 @@ void launch_variant(const Tensor& x, const Weight& weight, Tensor& qkv, Tensor&
static_assert((8192 % Schedule::BM) == 0 && (4096 % Schedule::BM) == 0);
const Output output{static_cast<__nv_bfloat16*>(qkv.data), static_cast<__nv_bfloat16*>(z.data)};
const dim3 grid(kRows / Schedule::BM, static_cast<unsigned>(div_up(x.ne[1], Schedule::BN)), 1u);
constexpr std::size_t kDynamicBytes = w8_rowsplit_gemm_mma_dynamic_bytes<Schedule>();
if constexpr (kDynamicBytes > kW8SmallTMmaStaticSharedBytes) {
static const cudaError_t attribute = cudaFuncSetAttribute(
w8_rowsplit_gemm_mma_kernel<Schedule, Full, W8Epilogue::Store, Output>,
cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(kDynamicBytes));
CUDA_CHECK(attribute);
}
w8_rowsplit_gemm_mma_kernel<Schedule, Full, W8Epilogue::Store, Output>
<<<grid, Schedule::THREADS, 0, stream>>>(static_cast<const __nv_bfloat16*>(x.data),
static_cast<const std::uint8_t*>(weight.qdata),
static_cast<const std::uint8_t*>(weight.scales),
output, kRows, kHidden, x.ne[1], kHidden);
<<<grid, Schedule::THREADS, kDynamicBytes, stream>>>(
static_cast<const __nv_bfloat16*>(x.data),
static_cast<const std::uint8_t*>(weight.qdata),
static_cast<const std::uint8_t*>(weight.scales), output, kRows, kHidden, x.ne[1],
kHidden);
}

} // namespace
Expand Down
22 changes: 20 additions & 2 deletions src/ops/gdn_input_proj/w8/w8_gdn_input_gemm_splitk.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -281,8 +281,17 @@ void launch_active_cols(const Tensor& x, const Weight& weight, Tensor& qkv, Tens
using Schedule = W8SmallTMmaDefaultSchedule<TileCols, ActiveCols>;
static_assert((8192 % kRowsPerCta) == 0 && (4096 % kRowsPerCta) == 0);
const Output output{static_cast<__nv_bfloat16*>(qkv.data), static_cast<__nv_bfloat16*>(z.data)};
constexpr std::size_t kDynamicBytes =
w8_small_t_mma_dynamic_bytes<Schedule, false, ActiveCols>();
if constexpr (kDynamicBytes > kW8SmallTMmaStaticSharedBytes) {
static const cudaError_t attribute = cudaFuncSetAttribute(
w8_small_t_mma_kernel<Geometry, ActiveCols, Schedule, Output,
W8SmallTMmaStoreEpilogue, W8SmallTMmaIdentityRows>,
cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(kDynamicBytes));
CUDA_CHECK(attribute);
}
w8_small_t_mma_kernel<Geometry, ActiveCols, Schedule>
<<<kRows / kRowsPerCta, Schedule::kThreads, 0, stream>>>(
<<<kRows / kRowsPerCta, Schedule::kThreads, kDynamicBytes, stream>>>(
static_cast<const __nv_bfloat16*>(x.data),
static_cast<const std::uint8_t*>(weight.qdata),
static_cast<const std::uint8_t*>(weight.scales), output);
Expand Down Expand Up @@ -320,8 +329,17 @@ void launch_active_cols_conv(const Tensor& x, const Weight& weight, const Tensor
},
static_cast<__nv_bfloat16*>(z.data),
};
constexpr std::size_t kDynamicBytes =
w8_small_t_mma_dynamic_bytes<Schedule, false, ActiveCols>();
if constexpr (kDynamicBytes > kW8SmallTMmaStaticSharedBytes) {
static const cudaError_t attribute = cudaFuncSetAttribute(
w8_small_t_mma_kernel<Geometry, ActiveCols, Schedule, Output,
W8GdnSplitKConvEpilogue<Publish>, W8SmallTMmaIdentityRows>,
cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(kDynamicBytes));
CUDA_CHECK(attribute);
}
w8_small_t_mma_kernel<Geometry, ActiveCols, Schedule, Output, W8GdnSplitKConvEpilogue<Publish>>
<<<kRows / kRowsPerCta, Schedule::kThreads, 0, stream>>>(
<<<kRows / kRowsPerCta, Schedule::kThreads, kDynamicBytes, stream>>>(
static_cast<const __nv_bfloat16*>(x.data),
static_cast<const std::uint8_t*>(weight.qdata),
static_cast<const std::uint8_t*>(weight.scales), ignored_output, epilogue);
Expand Down
16 changes: 12 additions & 4 deletions src/ops/linear/nvfp4/nvfp4_w4a4.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,10 +29,18 @@ void launch_gemm(const Weight& weight, Tensor& out, Nvfp4W4a4Workspace workspace
const Nvfp4ContiguousOutput output{static_cast<__nv_bfloat16*>(out.data),
Geometry::kOutputRows};
const float alpha = 1.0F / (weight.input_scale_divisor * weight.weight_scale_divisor);
nvfp4_w4a4_mma_kernel<Geometry, Schedule><<<grid, Schedule::kThreads, 0, stream>>>(
activation, static_cast<const std::uint8_t*>(weight.qdata),
static_cast<const std::uint8_t*>(weight.scales), tokens, alpha, Nvfp4IdentityEpilogue{},
output);
constexpr std::size_t kDynamicBytes = nvfp4_w4a4_mma_dynamic_bytes<Schedule>();
if constexpr (kDynamicBytes > kNvfp4W4a4StaticSharedBytes) {
static const cudaError_t attribute = cudaFuncSetAttribute(
nvfp4_w4a4_mma_kernel<Geometry, Schedule, Nvfp4IdentityEpilogue, Nvfp4ContiguousOutput>,
cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(kDynamicBytes));
CUDA_CHECK(attribute);
}
nvfp4_w4a4_mma_kernel<Geometry, Schedule>
<<<grid, Schedule::kThreads, kDynamicBytes, stream>>>(
activation, static_cast<const std::uint8_t*>(weight.qdata),
static_cast<const std::uint8_t*>(weight.scales), tokens, alpha,
Nvfp4IdentityEpilogue{}, output);
CUDA_CHECK(cudaGetLastError());
}

Expand Down
Loading