Skip to content

CUDA: MMQ tile configuration and FP4 accumulation for GB10 (sm_121) #198

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
danielhanchen wants to merge 5 commits into
base: master
Choose a base branch
from
+108 −5
Draft

CUDA: MMQ tile configuration and FP4 accumulation for GB10 (sm_121) #198

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
485088b
CUDA: MMQ tile configuration and FP4 accumulation for GB10 (sm_121)
danielhanchen Sep 6, 2026
92db8b7
CUDA: MMQ, carry the fallback tile choice from the caller instead of …
danielhanchen Sep 6, 2026
585f9c9
cuda: trim comments in the MMQ tile configuration change
danielhanchen Sep 9, 2026
d0a262e
CUDA: keep the GB10 MMQ tiles off every other Blackwell part
danielhanchen Sep 9, 2026
daccdb2
CUDA: keep the GB10 FP4 accumulator off every other Blackwell part
danielhanchen Sep 9, 2026
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
70 changes: 70 additions & 0 deletions ggml/src/ggml-cuda/mmq-config-dgx-spark.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,70 @@
// GB10 (sm_121a) ONLY. These tiles were swept on a DGX Spark, which is one SM cluster against
// unified LPDDR5X, and nothing here was measured on a discrete Blackwell part: an RTX 50-series
// card or an RTX PRO 6000 has its own SM count, its own GDDR7 bandwidth and its own occupancy
// limits, so a config tuned for the Spark is a guess there. Anything this table does not name
// falls through to the stock Blackwell config, so the override is exactly the swept entries.
static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_dgx_spark(ggml_type type, int J, bool fallback) {
CASE(GGML_TYPE_Q4_K, 128, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_Q4_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_Q4_K, 128, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_Q4_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_Q4_K, 128, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_Q4_K, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_Q4_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_Q4_K, 128, 2, 64, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_Q4_K, 128, 2, 64, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_Q4_K, 128, 2, 64, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_Q4_K, 128, 2, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);

CASE(GGML_TYPE_Q5_K, 128, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_Q5_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_Q5_K, 128, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_Q5_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_Q5_K, 128, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_Q5_K, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_Q5_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_Q5_K, 128, 2, 64, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_Q5_K, 128, 2, 64, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_Q5_K, 128, 2, 64, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_Q5_K, 128, 2, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);

CASE(GGML_TYPE_Q6_K, 128, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_Q6_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_Q6_K, 128, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_Q6_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_Q6_K, 128, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_Q6_K, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_Q6_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_Q6_K, 128, 2, 64, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_Q6_K, 128, 2, 64, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_Q6_K, 128, 2, 64, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_Q6_K, 128, 2, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);


CASE(GGML_TYPE_MXFP4, 128, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
CASE(GGML_TYPE_MXFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
CASE(GGML_TYPE_MXFP4, 128, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
CASE(GGML_TYPE_MXFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
CASE(GGML_TYPE_MXFP4, 128, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
CASE(GGML_TYPE_MXFP4, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
CASE(GGML_TYPE_MXFP4, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
CASE(GGML_TYPE_MXFP4, 128, 2, 64, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
CASE(GGML_TYPE_MXFP4, 128, 2, 64, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
CASE(GGML_TYPE_MXFP4, 128, 2, 64, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
CASE(GGML_TYPE_MXFP4, 128, 2, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);


CASE(GGML_TYPE_NVFP4, 128, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
CASE(GGML_TYPE_NVFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
CASE(GGML_TYPE_NVFP4, 128, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
CASE(GGML_TYPE_NVFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
CASE(GGML_TYPE_NVFP4, 128, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
CASE(GGML_TYPE_NVFP4, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
CASE(GGML_TYPE_NVFP4, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
CASE(GGML_TYPE_NVFP4, 128, 2, 64, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
CASE(GGML_TYPE_NVFP4, 128, 2, 64, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
CASE(GGML_TYPE_NVFP4, 128, 2, 64, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
CASE(GGML_TYPE_NVFP4, 128, 2, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);

return ggml_cuda_mmq_get_config_blackwell(type, J, fallback);
}
14 changes: 14 additions & 0 deletions ggml/src/ggml-cuda/mmq-vec-dot.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1236,6 +1236,19 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_

#pragma unroll
for (int n = 0; n < ntx; ++n) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK
// GB10 only, like the tile table this PR adds. Block-scaled MMA takes C as an
// accumulator input, so accumulating into the running sum removes a zeroed temporary
// per instruction, 805 FADD and all of the register spilling. It was swept on a DGX
// Spark and nowhere else, and PTX leaves mma's accumulation order unspecified, so it
// stays off every other Blackwell part for the same reason the tiles do.
static_assert(sizeof(tile_C) == tile_C::ne * sizeof(float), "tile_C must be a plain float array");
tile_C & C = *reinterpret_cast<tile_C *>(sum + (j0 / tile_C::J + n) * tile_C::ne);
#pragma unroll
for (int frag = 0; frag < nfrags; ++frag) {
mma_block_scaled_fp4<type>(C, A[n][frag], B[frag], scaleA[n][frag], scaleB[frag]);
}
#else
#pragma unroll
for (int frag = 0; frag < nfrags; ++frag) {
tile_C C = {};
Expand All @@ -1245,6 +1258,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
sum[(j0 / tile_C::J + n) * tile_C::ne + l] += C.x[l];
}
}
#endif // __CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK
}
}
}
Expand Down
12 changes: 9 additions & 3 deletions ggml/src/ggml-cuda/mmq.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -126,7 +126,13 @@ void ggml_cuda_mul_mat_q(
const int64_t s03 = src0->nb[3] / ts_src0;
const int64_t s3 = dst->nb[3] / ts_dst;

const bool fallback = ne01 % 128 != 0;
// GB10 only, matching the tile table it belongs to: mul_mat_id takes the fallback entries to
// keep the stock 128-row tile, because J comes from ncols_max but an expert holds only
// ncols_max*n_expert_used/n_experts of it, so a 64-row tile just doubles the row tiles paying
// for a mostly empty J. fallback only adds the src0->ne[1] range check.
const bool spark = GGML_CUDA_CC_IS_NVIDIA(cc)
&& ggml_cuda_highest_compiled_arch(cc) == GGML_CUDA_CC_DGX_SPARK;
const bool fallback = ne01 % 128 != 0 || (ids && spark);

const bool use_native_fp4 = blackwell_mma_available(cc) && (src0->type == GGML_TYPE_MXFP4 || src0->type == GGML_TYPE_NVFP4);
const size_t y_block_size = use_native_fp4 ? sizeof(block_fp4_mmq) : sizeof(block_q8_1_mmq);
Expand Down Expand Up @@ -171,7 +177,7 @@ void ggml_cuda_mul_mat_q(
ne00, ne01, ne1, s01, ne11, s1,
ne02, ne12, s02, s12, s2,
ne03, ne13, s03, s13, s3,
ne1};
ne1, fallback};
ggml_cuda_mul_mat_q_switch_type(ctx, args, stream);
return;
}
Expand Down Expand Up @@ -251,7 +257,7 @@ void ggml_cuda_mul_mat_q(
ne00, ne01, ne_get_rows, s01, ne_get_rows, s1,
ne02, ne02, s02, s12, s2,
ne03, ne13, s03, s13, s3,
ne12};
ne12, fallback};

ggml_cuda_mul_mat_q_switch_type(ctx, args, stream);
}
Expand Down
17 changes: 15 additions & 2 deletions ggml/src/ggml-cuda/mmq.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -216,6 +216,7 @@ struct ggml_cuda_mmq_config {
#include "mmq-config-pascal.cuh"
#include "mmq-config-ampere.cuh"
#include "mmq-config-blackwell.cuh"
#include "mmq-config-dgx-spark.cuh"

#include "mmq-config-cdna.cuh"
#include "mmq-config-rdna2.cuh"
Expand All @@ -241,6 +242,12 @@ static __host__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(const ggml_type ty
}
return ggml_cuda_mmq_get_config_rdna2(type, J, fallback);
}
// Ahead of the Blackwell branch, and an exact match rather than a range: the swept tiles
// below are a GB10 measurement and every other sm_12x part has its own SM count and memory.
// highest_compiled_arch, like mmvq.cu, so the host picks the table the device was built with.
if (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) == GGML_CUDA_CC_DGX_SPARK) {
return ggml_cuda_mmq_get_config_dgx_spark(type, J, fallback);
}
if (blackwell_mma_available(cc)) {
return ggml_cuda_mmq_get_config_blackwell(type, J, fallback);
}
Expand All @@ -264,7 +271,11 @@ static constexpr __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(ggml_t
return ggml_cuda_mmq_get_config_rdna2(type, J, fallback);
#endif // CDNA
#else
#ifdef BLACKWELL_MMA_AVAILABLE
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK
// Must mirror the host selector above, or the host would size shared memory for one config
// while the kernel compiled another.
return ggml_cuda_mmq_get_config_dgx_spark(type, J, fallback);
#elif defined(BLACKWELL_MMA_AVAILABLE)
return ggml_cuda_mmq_get_config_blackwell(type, J, fallback);
#elif __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA
return ggml_cuda_mmq_get_config_ampere(type, J, fallback);
Expand Down Expand Up @@ -1375,6 +1386,8 @@ struct mmq_args {
int64_t nchannels_x; int64_t nchannels_y; int64_t stride_channel_x; int64_t stride_channel_y; int64_t stride_channel_dst;
int64_t nsamples_x; int64_t nsamples_y; int64_t stride_sample_x; int64_t stride_sample_y; int64_t stride_sample_dst;
int64_t ncols_max;
// set by ggml_cuda_mul_mat_q, not re-derived, so it matches the config src1 padding was sized for
bool fallback;
};

static size_t mmq_get_nbytes_shared(const ggml_cuda_mmq_config & config, const int cc) {
Expand Down Expand Up @@ -1551,7 +1564,7 @@ void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args,

template <ggml_type type>
void mul_mat_q_case(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
if (args.nrows_x % 128 == 0) {
if (!args.fallback) {
constexpr bool fallback = false;
mul_mat_q_switch_J<type, fallback>(ctx, args, stream);
} else {
Expand Down