Skip to content

Optional BF16 for CUDA platform (Original PR #173) #205

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

Merged
guoqingbao merged 1 commit into from
Jul 27, 2025
Merged

Optional BF16 for CUDA platform (Original PR #173) #205

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
20 changes: 20 additions & 0 deletions kernels/src/marlin/marlin_dtypes.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,26 @@ class ScalarType<nv_bfloat16> {
static __host__ __device__ nv_bfloat16 inline float2num(const float x) {
return __float2bfloat16(x);
}

#else

static __device__ float inline num2float(const nv_bfloat16 x) {
return 0.0f;
}

static __device__ nv_bfloat162 inline num2num2(const nv_bfloat16 x) {
return nv_bfloat162();
}

static __device__ nv_bfloat162 inline nums2num2(const nv_bfloat16 x1,
const nv_bfloat16 x2) {
return nv_bfloat162();
}

static __host__ __device__ nv_bfloat16 inline float2num(const float x) {
return nv_bfloat16();
}

#endif
};

Expand Down
17 changes: 14 additions & 3 deletions kernels/src/marlin_cuda_kernel.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,14 +40,17 @@ __device__ inline void mma(const typename ScalarType<scalar_t>::FragA& a_frag,
: "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
: "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
"f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
} else if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
}
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
else if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
asm volatile(
"mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
"{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
: "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
: "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
"f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
}
#endif
}

// Instruction for loading a full 16x16 matrix fragment of operand A from shared
Expand Down Expand Up @@ -122,13 +125,14 @@ __device__ inline typename ScalarType<nv_bfloat16>::FragB
typename ScalarType<nv_bfloat16>::FragB frag_b;
static constexpr uint32_t MUL = 0x3F803F80;
static constexpr uint32_t ADD = 0xC308C308;

#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
frag_b[0] = __hfma2(*reinterpret_cast<nv_bfloat162*>(&lo),
*reinterpret_cast<const nv_bfloat162*>(&MUL),
*reinterpret_cast<const nv_bfloat162*>(&ADD));
frag_b[1] = __hfma2(*reinterpret_cast<nv_bfloat162*>(&hi),
*reinterpret_cast<const nv_bfloat162*>(&MUL),
*reinterpret_cast<const nv_bfloat162*>(&ADD));
#endif
return frag_b;
}

Expand Down Expand Up @@ -170,13 +174,14 @@ dequant<nv_bfloat16, ScalarTypeID::kU4>(int q) {
typename ScalarType<nv_bfloat16>::FragB frag_b;
static constexpr uint32_t MUL = 0x3F803F80;
static constexpr uint32_t ADD = 0xC300C300;

#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
frag_b[0] = __hfma2(*reinterpret_cast<nv_bfloat162*>(&lo),
*reinterpret_cast<const nv_bfloat162*>(&MUL),
*reinterpret_cast<const nv_bfloat162*>(&ADD));
frag_b[1] = __hfma2(*reinterpret_cast<nv_bfloat162*>(&hi),
*reinterpret_cast<const nv_bfloat162*>(&MUL),
*reinterpret_cast<const nv_bfloat162*>(&ADD));
#endif
return frag_b;
}
// Multiply dequantized values by the corresponding quantization scale; used
Expand Down Expand Up @@ -267,6 +272,12 @@ __global__ void Marlin(
// ensures good utilization of all SMs for many kinds of shape and GPU
// configurations, while requiring as few slow global cross-threadblock
// reductions as possible.
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
return;
}
#endif

using Dtype = ScalarType<scalar_t>;
using scalar_t2 = typename ScalarType<scalar_t>::scalar_t2;
using FragA = typename ScalarType<scalar_t>::FragA;
Expand Down
11 changes: 11 additions & 0 deletions kernels/src/pagedattention.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -430,6 +430,12 @@ __global__ void paged_attention_v1_kernel(
const int kv_block_stride,
const int kv_head_stride,
const float softscapping) {
if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
return;
#endif
}

paged_attention_kernel<scalar_t, HEAD_SIZE, BLOCK_SIZE, NUM_THREADS>(
/* exp_sums */ nullptr, /* max_logits */ nullptr,
out, q, k_cache, v_cache, num_kv_heads, scale, block_tables, context_lens,
Expand Down Expand Up @@ -460,6 +466,11 @@ __global__ void paged_attention_v2_kernel(
const int kv_block_stride,
const int kv_head_stride,
const float softscapping) {
if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
return;
#endif
}
paged_attention_kernel<scalar_t, HEAD_SIZE, BLOCK_SIZE, NUM_THREADS, PARTITION_SIZE>(
exp_sums, max_logits, tmp_out, q, k_cache, v_cache, num_kv_heads, scale,
block_tables, context_lens, max_num_blocks_per_seq, alibi_slopes,
Expand Down
58 changes: 50 additions & 8 deletions src/main.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -162,6 +162,55 @@ fn config_log(logger: ftail::Ftail, log_enable: bool, log_file: String) -> Resul
.map_err(candle_core::Error::wrap)
}

fn get_dtype(dtype: Option<String>) -> DType {
let dtype = match dtype.as_deref() {
Some("f16") => DType::F16,
Some("bf16") => DType::BF16,
Some("f32") => DType::F32,
Some(dtype) => panic!("Unsupported dtype {dtype}"),
None => DType::BF16,
};

#[cfg(feature = "cuda")]
let dtype = {
use candle_core::cuda_backend::cudarc::driver::result::{device, init};
use candle_core::cuda_backend::cudarc::driver::sys::CUdevice_attribute;
match (init(), device::get(0)) {
(Ok(_), Ok(d)) => {
let (compute_major, compute_minor) = unsafe {
(
device::get_attribute(
d,
CUdevice_attribute::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR,
)
.unwrap_or(8),
device::get_attribute(
d,
CUdevice_attribute::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR,
)
.unwrap_or(8),
)
};
info!(
"CUDA compute capability: {}.{}",
compute_major, compute_minor,
);
if dtype != DType::F32 && compute_major < 8 {
warn!(
"CUDA compute capability: {} (<8), switched to F16 cause no BF16 support.",
compute_major
);
DType::F16
} else {
dtype
}
}
_ => dtype,
}
};
dtype
}

#[tokio::main]
#[allow(unused_mut)]
async fn main() -> Result<()> {
Expand All @@ -180,14 +229,7 @@ async fn main() -> Result<()> {

let (paths, gguf) = loader.prepare_model_weights(args.hf_token, args.hf_token_path)?;

let dtype = match args.dtype.as_deref() {
Some("f16") => DType::F16,
Some("bf16") => DType::BF16,
Some("f32") => DType::F32,
Some(dtype) => panic!("Unsupported dtype {dtype}"),
None => DType::BF16,
};

let dtype = get_dtype(args.dtype);
let device_ids: Vec<usize> = match args.device_ids {
Some(ids) => ids,
_ => vec![0usize],
Expand Down
Loading