Skip to content

CUDA: dispatch the chunked gated_delta_net in CU mode when it fills the device #79

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

Open
roberteg16 wants to merge 1 commit into
base: gfx11
Choose a base branch
from
+45 −18
Open

CUDA: dispatch the chunked gated_delta_net in CU mode when it fills the device #79

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
63 changes: 45 additions & 18 deletions ggml/src/ggml-cuda/gated_delta_net_chunked.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -109,23 +109,34 @@ bool ggml_cuda_gdn_chunked_supported(bool kda, bool keep_rs, int64_t S_v, int64_
#define GDN_MIN_WAVES_PER_SIMD 1
#endif

__global__ void __launch_bounds__(GDN_BLOCK_SIZE, GDN_MIN_WAVES_PER_SIMD)
gdn_chunked_f32(const float * __restrict__ q,
const float * __restrict__ k,
const float * __restrict__ v,
const float * __restrict__ g,
const float * __restrict__ beta,
const float * __restrict__ state_in,
float * __restrict__ dst,
float * __restrict__ state_out,
const int64_t n_tokens,
const int64_t n_heads,
const int64_t sq1, const int64_t sq2, const int64_t sq3,
const int64_t sv1, const int64_t sv2, const int64_t sv3,
const int64_t sb1, const int64_t sb2, const int64_t sb3,
const uint3 neqk1_magic,
const uint3 rq3_magic,
const float scale) {
// CU mode confines a block to one CU, which halves the LDS contention domain. It costs no residency:
// a WGP holds floor(128 KB / 57.25 KB) = 2 of these blocks in WGP mode and 2*floor(64 KB / 57.25 KB)
// = 2 in CU mode. It only pays once every CU has a block to run, hence the launch-time choice below.
#if defined(GGML_USE_HIP) && defined(RDNA3_5)
#define GDN_CU_MODE __attribute__((target("cumode")))
#else
#define GDN_CU_MODE
#endif

// The work-group processor mode is a function attribute, so it cannot be a template parameter: the
// two entry points below share this body and differ only in that attribute.
#define GDN_CHUNKED_PARAMS \
const float * __restrict__ q, const float * __restrict__ k, \
const float * __restrict__ v, const float * __restrict__ g, \
const float * __restrict__ beta, const float * __restrict__ state_in, \
float * __restrict__ dst, float * __restrict__ state_out, \
const int64_t n_tokens, const int64_t n_heads, \
const int64_t sq1, const int64_t sq2, const int64_t sq3, \
const int64_t sv1, const int64_t sv2, const int64_t sv3, \
const int64_t sb1, const int64_t sb2, const int64_t sb3, \
const uint3 neqk1_magic, const uint3 rq3_magic, \
const float scale

#define GDN_CHUNKED_ARGS \
q, k, v, g, beta, state_in, dst, state_out, n_tokens, n_heads, \
sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale

static __device__ __forceinline__ void gdn_chunked_body(GDN_CHUNKED_PARAMS) {
constexpr int CHUNK = GDN_CHUNK;
constexpr int HEAD_DIM = GDN_HEAD_DIM;
constexpr int COL_LANES = GDN_COL_LANES;
Expand Down Expand Up @@ -399,6 +410,16 @@ gdn_chunked_f32(const float * __restrict__ q,
}
}

__global__ void __launch_bounds__(GDN_BLOCK_SIZE, GDN_MIN_WAVES_PER_SIMD)
gdn_chunked_f32_wgp(GDN_CHUNKED_PARAMS) {
gdn_chunked_body(GDN_CHUNKED_ARGS);
}

__global__ void GDN_CU_MODE __launch_bounds__(GDN_BLOCK_SIZE, GDN_MIN_WAVES_PER_SIMD)
gdn_chunked_f32_cu(GDN_CHUNKED_PARAMS) {
gdn_chunked_body(GDN_CHUNKED_ARGS);
}

void ggml_cuda_gdn_chunked(ggml_backend_cuda_context & ctx, const ggml_cuda_gdn_chunked_args & args) {
GGML_ASSERT(args.S_v == GDN_HEAD_DIM);

Expand All @@ -410,7 +431,13 @@ void ggml_cuda_gdn_chunked(ggml_backend_cuda_context & ctx, const ggml_cuda_gdn_

const ggml_cuda_kernel_launch_params launch_params(block_nums, block_dims, 0, ctx.stream());

ggml_cuda_kernel_launch(gdn_chunked_f32, launch_params,
// nsm counts WGPs on RDNA, and a block occupies one CU in CU mode but spreads over the whole WGP
// in WGP mode. Below one block per WGP, CU mode leaves the second CU of each WGP idle and costs
// about a quarter of the op; above it every CU is busy either way and CU mode is worth ~10%.
const bool cu_mode = args.H * args.n_seqs > ggml_cuda_info().devices[ggml_cuda_get_device()].nsm;
const auto kernel = cu_mode ? gdn_chunked_f32_cu : gdn_chunked_f32_wgp;

ggml_cuda_kernel_launch(kernel, launch_params,
args.q, args.k, args.v, args.g, args.beta, args.state_in, args.dst, args.state_out,
args.n_tokens, args.H,
args.sq1, args.sq2, args.sq3, args.sv1, args.sv2, args.sv3, args.sb1, args.sb2, args.sb3,
Expand Down
Loading