Skip to content

GitLab

"tests/vscode:/vscode.git/clone" did not exist on "89514f0541f312945854236132a5f4ea2516fa86"
Commit 1a8114bf authored by fengzch-das's avatar fengzch-das
Browse files

hipify code

parent c0177256
Pipeline #3049 canceled with stages
Hide whitespace changes
Inline Side-by-side
Showing with 254 additions and 248 deletions
nunchaku/csrc/gemm.h
View file @ 1a8114bf
...@@ -12,8 +12,8 @@ public: ...@@ -12,8 +12,8 @@ public:
spdlog::info("Initializing QuantizedGEMM"); spdlog::info("Initializing QuantizedGEMM");
size_t val = 0; size_t val = 0;
checkCUDA(cudaDeviceSetLimit(cudaLimitStackSize, 8192)); checkCUDA(hipDeviceSetLimit(hipLimitStackSize, 8192));
checkCUDA(cudaDeviceGetLimit(&val, cudaLimitStackSize)); checkCUDA(hipDeviceGetLimit(&val, hipLimitStackSize));
spdlog::debug("Stack={}", val); spdlog::debug("Stack={}", val);
net = std::make_unique<GEMM_W4A4>((int)in_features, net = std::make_unique<GEMM_W4A4>((int)in_features,
...@@ -42,7 +42,7 @@ public: ...@@ -42,7 +42,7 @@ public:
std::string dumpTensorBF16(Tensor x) { std::string dumpTensorBF16(Tensor x) {
std::stringstream ss; std::stringstream ss;
for (int i = 0; i < 256; i++) { for (int i = 0; i < 256; i++) {
ss << spdlog::fmt_lib::format("{:.3f} ", (float)(x.data_ptr<__nv_bfloat16>()[i])); ss << spdlog::fmt_lib::format("{:.3f} ", (float)(x.data_ptr<__hip_bfloat16>()[i]));
} }
ss << std::endl; ss << std::endl;
return ss.str(); return ss.str();
......
nunchaku/csrc/gemm88.h
View file @ 1a8114bf
...@@ -12,8 +12,8 @@ public: ...@@ -12,8 +12,8 @@ public:
spdlog::info("Initializing QuantizedGEMM88"); spdlog::info("Initializing QuantizedGEMM88");
size_t val = 0; size_t val = 0;
checkCUDA(cudaDeviceSetLimit(cudaLimitStackSize, 8192)); checkCUDA(hipDeviceSetLimit(hipLimitStackSize, 8192));
checkCUDA(cudaDeviceGetLimit(&val, cudaLimitStackSize)); checkCUDA(hipDeviceGetLimit(&val, hipLimitStackSize));
spdlog::debug("Stack={}", val); spdlog::debug("Stack={}", val);
net = std::make_unique<GEMM_W8A8>( net = std::make_unique<GEMM_W8A8>(
......
nunchaku/csrc/utils.h
View file @ 1a8114bf
...@@ -8,27 +8,27 @@ namespace nunchaku::utils { ...@@ -8,27 +8,27 @@ namespace nunchaku::utils {
void set_cuda_stack_limit(int64_t newval) { void set_cuda_stack_limit(int64_t newval) {
size_t val = 0; size_t val = 0;
checkCUDA(cudaDeviceSetLimit(cudaLimitStackSize, (size_t)newval)); checkCUDA(hipDeviceSetLimit(hipLimitStackSize, (size_t)newval));
checkCUDA(cudaDeviceGetLimit(&val, cudaLimitStackSize)); checkCUDA(hipDeviceGetLimit(&val, hipLimitStackSize));
spdlog::debug("Stack={}", val); spdlog::debug("Stack={}", val);
} }
void disable_memory_auto_release() { void disable_memory_auto_release() {
int device; int device;
checkCUDA(cudaGetDevice(&device)); checkCUDA(hipGetDevice(&device));
cudaMemPool_t mempool; hipMemPool_t mempool;
checkCUDA(cudaDeviceGetDefaultMemPool(&mempool, device)); checkCUDA(hipDeviceGetDefaultMemPool(&mempool, device));
uint64_t threshold = UINT64_MAX; uint64_t threshold = UINT64_MAX;
checkCUDA(cudaMemPoolSetAttribute(mempool, cudaMemPoolAttrReleaseThreshold, &threshold)); checkCUDA(hipMemPoolSetAttribute(mempool, hipMemPoolAttrReleaseThreshold, &threshold));
} }
void trim_memory() { void trim_memory() {
int device; int device;
checkCUDA(cudaGetDevice(&device)); checkCUDA(hipGetDevice(&device));
cudaMemPool_t mempool; hipMemPool_t mempool;
checkCUDA(cudaDeviceGetDefaultMemPool(&mempool, device)); checkCUDA(hipDeviceGetDefaultMemPool(&mempool, device));
size_t bytesToKeep = 0; size_t bytesToKeep = 0;
checkCUDA(cudaMemPoolTrimTo(mempool, bytesToKeep)); checkCUDA(hipMemPoolTrimTo(mempool, bytesToKeep));
} }
void set_faster_i2f_mode(std::string mode) { void set_faster_i2f_mode(std::string mode) {
......
src/FluxModel.cpp
View file @ 1a8114bf
...@@ -4,7 +4,7 @@ ...@@ -4,7 +4,7 @@
#include "kernels/zgemm/zgemm.h" #include "kernels/zgemm/zgemm.h"
#include "flash_api.h" #include "flash_api.h"
#include "activation.h" #include "activation.h"
#include <nvtx3/nvToolsExt.h> #include <nvtx3/roctracer/roctx.h>
#include <pybind11/functional.h> #include <pybind11/functional.h>
...@@ -296,7 +296,7 @@ FluxSingleTransformerBlock::FluxSingleTransformerBlock(int dim, ...@@ -296,7 +296,7 @@ FluxSingleTransformerBlock::FluxSingleTransformerBlock(int dim,
Tensor FluxSingleTransformerBlock::forward(Tensor hidden_states, Tensor temb, Tensor rotary_emb) { Tensor FluxSingleTransformerBlock::forward(Tensor hidden_states, Tensor temb, Tensor rotary_emb) {
nvtxRangePushA("FluxSingleTransformerBlock"); roctxRangePushA("FluxSingleTransformerBlock");
const int batch_size = hidden_states.shape[0]; const int batch_size = hidden_states.shape[0];
const int num_tokens = hidden_states.shape[1]; const int num_tokens = hidden_states.shape[1];
...@@ -370,14 +370,14 @@ Tensor FluxSingleTransformerBlock::forward(Tensor hidden_states, Tensor temb, Te ...@@ -370,14 +370,14 @@ Tensor FluxSingleTransformerBlock::forward(Tensor hidden_states, Tensor temb, Te
attn_output = o.slice(1, 0, num_tokens); attn_output = o.slice(1, 0, num_tokens);
} else { } else {
attn_output = Tensor::allocate({batch_size, num_tokens, num_heads * dim_head}, o.scalar_type(), o.device()); attn_output = Tensor::allocate({batch_size, num_tokens, num_heads * dim_head}, o.scalar_type(), o.device());
checkCUDA(cudaMemcpy2DAsync(attn_output.data_ptr(), checkCUDA(hipMemcpy2DAsync(attn_output.data_ptr(),
attn_output.stride(0) * attn_output.scalar_size(), attn_output.stride(0) * attn_output.scalar_size(),
o.data_ptr(), o.data_ptr(),
o.stride(0) * o.scalar_size(), o.stride(0) * o.scalar_size(),
attn_output.stride(0) * attn_output.scalar_size(), attn_output.stride(0) * attn_output.scalar_size(),
batch_size, batch_size,
cudaMemcpyDeviceToDevice, hipMemcpyDeviceToDevice,
getCurrentCUDAStream())); getCurrentHIPStreamMasqueradingAsCUDA()));
} }
} else { } else {
assert(false); assert(false);
...@@ -397,7 +397,7 @@ Tensor FluxSingleTransformerBlock::forward(Tensor hidden_states, Tensor temb, Te ...@@ -397,7 +397,7 @@ Tensor FluxSingleTransformerBlock::forward(Tensor hidden_states, Tensor temb, Te
// kernels::mul_add(hidden_states, gate, residual); // kernels::mul_add(hidden_states, gate, residual);
kernels::mul_add_batch(hidden_states, gate, true, 0.0, residual, true); kernels::mul_add_batch(hidden_states, gate, true, 0.0, residual, true);
nvtxRangePop(); roctxRangePop();
return hidden_states; return hidden_states;
} }
...@@ -439,9 +439,9 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states, ...@@ -439,9 +439,9 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states,
int batch_size = hidden_states.shape[0]; int batch_size = hidden_states.shape[0];
assert(encoder_hidden_states.shape[0] == batch_size); assert(encoder_hidden_states.shape[0] == batch_size);
nvtxRangePushA("JointTransformerBlock"); roctxRangePushA("JointTransformerBlock");
nvtxRangePushA("AdaNorm"); roctxRangePushA("AdaNorm");
int num_tokens_img = hidden_states.shape[1]; int num_tokens_img = hidden_states.shape[1];
int num_tokens_txt = encoder_hidden_states.shape[1]; int num_tokens_txt = encoder_hidden_states.shape[1];
...@@ -468,9 +468,9 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states, ...@@ -468,9 +468,9 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states,
constexpr int POOL_SIZE = Attention::POOL_SIZE; constexpr int POOL_SIZE = Attention::POOL_SIZE;
nvtxRangePop(); roctxRangePop();
auto stream = getCurrentCUDAStream(); auto stream = getCurrentHIPStreamMasqueradingAsCUDA();
int num_tokens_img_pad = 0, num_tokens_txt_pad = 0; int num_tokens_img_pad = 0, num_tokens_txt_pad = 0;
Tensor raw_attn_output; Tensor raw_attn_output;
...@@ -483,7 +483,7 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states, ...@@ -483,7 +483,7 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states,
Tensor pool; Tensor pool;
{ {
nvtxRangePushA("qkv_proj"); roctxRangePushA("qkv_proj");
const bool blockSparse = sparsityRatio > 0; const bool blockSparse = sparsityRatio > 0;
...@@ -535,7 +535,7 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states, ...@@ -535,7 +535,7 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states,
debug("qkv_context", qkv_context); debug("qkv_context", qkv_context);
} }
nvtxRangePop(); roctxRangePop();
} }
spdlog::debug("concat={}", concat.shape.str()); spdlog::debug("concat={}", concat.shape.str());
...@@ -543,7 +543,7 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states, ...@@ -543,7 +543,7 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states,
assert(concat.shape[2] == num_heads * dim_head * 3); assert(concat.shape[2] == num_heads * dim_head * 3);
nvtxRangePushA("Attention"); roctxRangePushA("Attention");
if (pool.valid()) { if (pool.valid()) {
raw_attn_output = attn.forward(concat, pool, sparsityRatio); raw_attn_output = attn.forward(concat, pool, sparsityRatio);
...@@ -551,7 +551,7 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states, ...@@ -551,7 +551,7 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states,
raw_attn_output = attn.forward(concat); raw_attn_output = attn.forward(concat);
} }
nvtxRangePop(); roctxRangePop();
spdlog::debug("raw_attn_output={}", raw_attn_output.shape.str()); spdlog::debug("raw_attn_output={}", raw_attn_output.shape.str());
...@@ -564,7 +564,7 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states, ...@@ -564,7 +564,7 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states,
Tensor concat_q, concat_k, concat_v; Tensor concat_q, concat_k, concat_v;
{ {
nvtxRangePushA("qkv_proj"); roctxRangePushA("qkv_proj");
concat_q = Tensor::allocate({batch_size, num_heads, num_tokens_img_pad + num_tokens_txt_pad, dim_head}, concat_q = Tensor::allocate({batch_size, num_heads, num_tokens_img_pad + num_tokens_txt_pad, dim_head},
Tensor::FP16, Tensor::FP16,
...@@ -606,18 +606,18 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states, ...@@ -606,18 +606,18 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states,
debug("concat_k", concat_k); debug("concat_k", concat_k);
debug("concat_v", concat_v); debug("concat_v", concat_v);
nvtxRangePop(); roctxRangePop();
} }
raw_attn_output = Tensor::allocate({batch_size, num_tokens_img_pad + num_tokens_txt_pad, num_heads * dim_head}, raw_attn_output = Tensor::allocate({batch_size, num_tokens_img_pad + num_tokens_txt_pad, num_heads * dim_head},
norm1_output.x.scalar_type(), norm1_output.x.scalar_type(),
norm1_output.x.device()); norm1_output.x.device());
nvtxRangePushA("Attention"); roctxRangePushA("Attention");
kernels::attention_fp16(concat_q, concat_k, concat_v, raw_attn_output, pow(dim_head, (-0.5))); kernels::attention_fp16(concat_q, concat_k, concat_v, raw_attn_output, pow(dim_head, (-0.5)));
nvtxRangePop(); roctxRangePop();
raw_attn_output = raw_attn_output =
raw_attn_output.view({batch_size, num_tokens_img_pad + num_tokens_txt_pad, num_heads, dim_head}); raw_attn_output.view({batch_size, num_tokens_img_pad + num_tokens_txt_pad, num_heads, dim_head});
...@@ -628,7 +628,7 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states, ...@@ -628,7 +628,7 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states,
debug("raw_attn_output", raw_attn_output); debug("raw_attn_output", raw_attn_output);
{ {
nvtxRangePushA("o_proj"); roctxRangePushA("o_proj");
auto &&[_, gate_msa, shift_mlp, scale_mlp, gate_mlp] = norm1_output; auto &&[_, gate_msa, shift_mlp, scale_mlp, gate_mlp] = norm1_output;
...@@ -642,14 +642,14 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states, ...@@ -642,14 +642,14 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states,
raw_attn_output_split = Tensor::allocate({batch_size, num_tokens_img, num_heads * dim_head}, raw_attn_output_split = Tensor::allocate({batch_size, num_tokens_img, num_heads * dim_head},
raw_attn_output.scalar_type(), raw_attn_output.scalar_type(),
raw_attn_output.device()); raw_attn_output.device());
checkCUDA(cudaMemcpy2DAsync(raw_attn_output_split.data_ptr(), checkCUDA(hipMemcpy2DAsync(raw_attn_output_split.data_ptr(),
num_tokens_img * num_heads * dim_head * raw_attn_output_split.scalar_size(), num_tokens_img * num_heads * dim_head * raw_attn_output_split.scalar_size(),
raw_attn_output.data_ptr(), raw_attn_output.data_ptr(),
(num_tokens_img_pad + num_tokens_txt_pad) * num_heads * dim_head * (num_tokens_img_pad + num_tokens_txt_pad) * num_heads * dim_head *
raw_attn_output.scalar_size(), raw_attn_output.scalar_size(),
num_tokens_img * num_heads * dim_head * raw_attn_output_split.scalar_size(), num_tokens_img * num_heads * dim_head * raw_attn_output_split.scalar_size(),
batch_size, batch_size,
cudaMemcpyDeviceToDevice, hipMemcpyDeviceToDevice,
stream)); stream));
} }
...@@ -665,8 +665,8 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states, ...@@ -665,8 +665,8 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states,
kernels::mul_add_batch(attn_output, gate_msa, true, 0.0, hidden_states, true); kernels::mul_add_batch(attn_output, gate_msa, true, 0.0, hidden_states, true);
hidden_states = std::move(attn_output); hidden_states = std::move(attn_output);
nvtxRangePop(); roctxRangePop();
nvtxRangePushA("MLP"); roctxRangePushA("MLP");
spdlog::debug("attn_output={}", hidden_states.shape.str()); spdlog::debug("attn_output={}", hidden_states.shape.str());
...@@ -691,7 +691,7 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states, ...@@ -691,7 +691,7 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states,
kernels::mul_add_batch(ff_output, gate_mlp, true, 0.0, hidden_states, true); kernels::mul_add_batch(ff_output, gate_mlp, true, 0.0, hidden_states, true);
hidden_states = std::move(ff_output); hidden_states = std::move(ff_output);
nvtxRangePop(); roctxRangePop();
spdlog::debug("ff_output={}", hidden_states.shape.str()); spdlog::debug("ff_output={}", hidden_states.shape.str());
} }
...@@ -701,7 +701,7 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states, ...@@ -701,7 +701,7 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states,
} }
{ {
nvtxRangePushA("o_proj_context"); roctxRangePushA("o_proj_context");
auto &&[_, gate_msa, shift_mlp, scale_mlp, gate_mlp] = norm1_context_output; auto &&[_, gate_msa, shift_mlp, scale_mlp, gate_mlp] = norm1_context_output;
...@@ -713,7 +713,7 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states, ...@@ -713,7 +713,7 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states,
raw_attn_output_split = Tensor::allocate({batch_size, num_tokens_txt, num_heads * dim_head}, raw_attn_output_split = Tensor::allocate({batch_size, num_tokens_txt, num_heads * dim_head},
raw_attn_output.scalar_type(), raw_attn_output.scalar_type(),
raw_attn_output.device()); raw_attn_output.device());
checkCUDA(cudaMemcpy2DAsync(raw_attn_output_split.data_ptr(), checkCUDA(hipMemcpy2DAsync(raw_attn_output_split.data_ptr(),
num_tokens_txt * num_heads * dim_head * raw_attn_output_split.scalar_size(), num_tokens_txt * num_heads * dim_head * raw_attn_output_split.scalar_size(),
raw_attn_output.data_ptr<char>() + num_tokens_img_pad * num_heads * dim_head * raw_attn_output.data_ptr<char>() + num_tokens_img_pad * num_heads * dim_head *
raw_attn_output_split.scalar_size(), raw_attn_output_split.scalar_size(),
...@@ -721,7 +721,7 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states, ...@@ -721,7 +721,7 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states,
raw_attn_output.scalar_size(), raw_attn_output.scalar_size(),
num_tokens_txt * num_heads * dim_head * raw_attn_output_split.scalar_size(), num_tokens_txt * num_heads * dim_head * raw_attn_output_split.scalar_size(),
batch_size, batch_size,
cudaMemcpyDeviceToDevice, hipMemcpyDeviceToDevice,
stream)); stream));
} }
...@@ -738,8 +738,8 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states, ...@@ -738,8 +738,8 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states,
kernels::mul_add_batch(attn_output, gate_msa, true, 0.0, encoder_hidden_states, true); kernels::mul_add_batch(attn_output, gate_msa, true, 0.0, encoder_hidden_states, true);
encoder_hidden_states = std::move(attn_output); encoder_hidden_states = std::move(attn_output);
nvtxRangePop(); roctxRangePop();
nvtxRangePushA("MLP"); roctxRangePushA("MLP");
spdlog::debug("attn_output={}", encoder_hidden_states.shape.str()); spdlog::debug("attn_output={}", encoder_hidden_states.shape.str());
...@@ -766,12 +766,12 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states, ...@@ -766,12 +766,12 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states,
kernels::mul_add_batch(ff_output, gate_mlp, true, 0.0, encoder_hidden_states, true); kernels::mul_add_batch(ff_output, gate_mlp, true, 0.0, encoder_hidden_states, true);
encoder_hidden_states = std::move(ff_output); encoder_hidden_states = std::move(ff_output);
nvtxRangePop(); roctxRangePop();
spdlog::debug("ff_output={}", encoder_hidden_states.shape.str()); spdlog::debug("ff_output={}", encoder_hidden_states.shape.str());
} }
nvtxRangePop(); roctxRangePop();
return {hidden_states, encoder_hidden_states}; return {hidden_states, encoder_hidden_states};
} }
...@@ -831,12 +831,12 @@ Tensor JointTransformerBlock::get_q_heads(Tensor hidden_states, ...@@ -831,12 +831,12 @@ Tensor JointTransformerBlock::get_q_heads(Tensor hidden_states,
size_t width = C * E; size_t width = C * E;
size_t height = R; size_t height = R;
Tensor out = Tensor::allocate({B, R, C}, t.scalarType, t.device()); Tensor out = Tensor::allocate({B, R, C}, t.scalarType, t.device());
auto stream = getCurrentCUDAStream(); auto stream = getCurrentHIPStreamMasqueradingAsCUDA();
for (int b = 0; b < B; ++b) { for (int b = 0; b < B; ++b) {
const void *src = (const char *)t.data_ptr<char>() + t.stride(0) * b * E; const void *src = (const char *)t.data_ptr<char>() + t.stride(0) * b * E;
void *dst = (char *)out.data_ptr<char>() + out.stride(0) * b * E; void *dst = (char *)out.data_ptr<char>() + out.stride(0) * b * E;
checkCUDA( checkCUDA(
cudaMemcpy2DAsync(dst, dst_pitch, src, src_pitch, width, height, cudaMemcpyDeviceToDevice, stream)); hipMemcpy2DAsync(dst, dst_pitch, src, src_pitch, width, height, hipMemcpyDeviceToDevice, stream));
} }
return out; return out;
}; };
...@@ -852,9 +852,9 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra ...@@ -852,9 +852,9 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra
int batch_size = hidden_states.shape[0]; int batch_size = hidden_states.shape[0];
assert(encoder_hidden_states.shape[0] == batch_size); assert(encoder_hidden_states.shape[0] == batch_size);
nvtxRangePushA("JointTransformerBlock"); roctxRangePushA("JointTransformerBlock");
nvtxRangePushA("AdaNorm"); roctxRangePushA("AdaNorm");
int num_tokens_img = hidden_states.shape[1]; int num_tokens_img = hidden_states.shape[1];
int num_tokens_txt = encoder_hidden_states.shape[1]; int num_tokens_txt = encoder_hidden_states.shape[1];
...@@ -877,12 +877,12 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra ...@@ -877,12 +877,12 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra
Tensor out = Tensor::allocate({B, R, C}, t.scalarType, t.device()); Tensor out = Tensor::allocate({B, R, C}, t.scalarType, t.device());
auto stream = getCurrentCUDAStream(); auto stream = getCurrentHIPStreamMasqueradingAsCUDA();
for (int b = 0; b < B; ++b) { for (int b = 0; b < B; ++b) {
const void *src = (const char *)t.data_ptr<char>() + t.stride(0) * b * E; const void *src = (const char *)t.data_ptr<char>() + t.stride(0) * b * E;
void *dst = (char *)out.data_ptr<char>() + out.stride(0) * b * E; void *dst = (char *)out.data_ptr<char>() + out.stride(0) * b * E;
checkCUDA( checkCUDA(
cudaMemcpy2DAsync(dst, dst_pitch, src, src_pitch, width, height, cudaMemcpyDeviceToDevice, stream)); hipMemcpy2DAsync(dst, dst_pitch, src, src_pitch, width, height, hipMemcpyDeviceToDevice, stream));
} }
return out; return out;
}; };
...@@ -906,9 +906,9 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra ...@@ -906,9 +906,9 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra
constexpr int POOL_SIZE = Attention::POOL_SIZE; constexpr int POOL_SIZE = Attention::POOL_SIZE;
nvtxRangePop(); roctxRangePop();
auto stream = getCurrentCUDAStream(); auto stream = getCurrentHIPStreamMasqueradingAsCUDA();
int num_tokens_img_pad = 0, num_tokens_txt_pad = 0; int num_tokens_img_pad = 0, num_tokens_txt_pad = 0;
Tensor raw_attn_output; Tensor raw_attn_output;
...@@ -921,7 +921,7 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra ...@@ -921,7 +921,7 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra
Tensor pool; Tensor pool;
{ {
nvtxRangePushA("qkv_proj"); roctxRangePushA("qkv_proj");
const bool blockSparse = sparsityRatio > 0; const bool blockSparse = sparsityRatio > 0;
...@@ -973,7 +973,7 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra ...@@ -973,7 +973,7 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra
debug("qkv_context", qkv_context); debug("qkv_context", qkv_context);
} }
nvtxRangePop(); roctxRangePop();
} }
spdlog::debug("concat={}", concat.shape.str()); spdlog::debug("concat={}", concat.shape.str());
...@@ -981,7 +981,7 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra ...@@ -981,7 +981,7 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra
assert(concat.shape[2] == num_heads * dim_head * 3); assert(concat.shape[2] == num_heads * dim_head * 3);
nvtxRangePushA("Attention"); roctxRangePushA("Attention");
if (pool.valid()) { if (pool.valid()) {
raw_attn_output = attn.forward(concat, pool, sparsityRatio); raw_attn_output = attn.forward(concat, pool, sparsityRatio);
...@@ -989,7 +989,7 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra ...@@ -989,7 +989,7 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra
raw_attn_output = attn.forward(concat); raw_attn_output = attn.forward(concat);
} }
nvtxRangePop(); roctxRangePop();
spdlog::debug("raw_attn_output={}", raw_attn_output.shape.str()); spdlog::debug("raw_attn_output={}", raw_attn_output.shape.str());
...@@ -1009,7 +1009,7 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra ...@@ -1009,7 +1009,7 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra
Tensor concat_q, concat_k, concat_v; Tensor concat_q, concat_k, concat_v;
{ {
nvtxRangePushA("qkv_proj"); roctxRangePushA("qkv_proj");
concat_q = Tensor::allocate({batch_size, num_heads, num_tokens_img_pad + num_tokens_txt_pad, dim_head}, concat_q = Tensor::allocate({batch_size, num_heads, num_tokens_img_pad + num_tokens_txt_pad, dim_head},
Tensor::FP16, Tensor::FP16,
...@@ -1051,18 +1051,18 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra ...@@ -1051,18 +1051,18 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra
debug("concat_k", concat_k); debug("concat_k", concat_k);
debug("concat_v", concat_v); debug("concat_v", concat_v);
nvtxRangePop(); roctxRangePop();
} }
raw_attn_output = Tensor::allocate({batch_size, num_tokens_img_pad + num_tokens_txt_pad, num_heads * dim_head}, raw_attn_output = Tensor::allocate({batch_size, num_tokens_img_pad + num_tokens_txt_pad, num_heads * dim_head},
norm1_output.x.scalar_type(), norm1_output.x.scalar_type(),
norm1_output.x.device()); norm1_output.x.device());
nvtxRangePushA("Attention"); roctxRangePushA("Attention");
kernels::attention_fp16(concat_q, concat_k, concat_v, raw_attn_output, pow(dim_head, (-0.5))); kernels::attention_fp16(concat_q, concat_k, concat_v, raw_attn_output, pow(dim_head, (-0.5)));
nvtxRangePop(); roctxRangePop();
raw_attn_output = raw_attn_output =
raw_attn_output.view({batch_size, num_tokens_img_pad + num_tokens_txt_pad, num_heads, dim_head}); raw_attn_output.view({batch_size, num_tokens_img_pad + num_tokens_txt_pad, num_heads, dim_head});
...@@ -1075,7 +1075,7 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra ...@@ -1075,7 +1075,7 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra
debug("raw_attn_output", raw_attn_output); debug("raw_attn_output", raw_attn_output);
{ {
nvtxRangePushA("o_proj"); roctxRangePushA("o_proj");
auto &&[_, gate_msa, shift_mlp, scale_mlp, gate_mlp] = norm1_output; auto &&[_, gate_msa, shift_mlp, scale_mlp, gate_mlp] = norm1_output;
...@@ -1089,14 +1089,14 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra ...@@ -1089,14 +1089,14 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra
raw_attn_output_split = Tensor::allocate({batch_size, num_tokens_img, num_heads * dim_head}, raw_attn_output_split = Tensor::allocate({batch_size, num_tokens_img, num_heads * dim_head},
raw_attn_output.scalar_type(), raw_attn_output.scalar_type(),
raw_attn_output.device()); raw_attn_output.device());
checkCUDA(cudaMemcpy2DAsync(raw_attn_output_split.data_ptr(), checkCUDA(hipMemcpy2DAsync(raw_attn_output_split.data_ptr(),
num_tokens_img * num_heads * dim_head * raw_attn_output_split.scalar_size(), num_tokens_img * num_heads * dim_head * raw_attn_output_split.scalar_size(),
raw_attn_output.data_ptr(), raw_attn_output.data_ptr(),
(num_tokens_img_pad + num_tokens_txt_pad) * num_heads * dim_head * (num_tokens_img_pad + num_tokens_txt_pad) * num_heads * dim_head *
raw_attn_output.scalar_size(), raw_attn_output.scalar_size(),
num_tokens_img * num_heads * dim_head * raw_attn_output_split.scalar_size(), num_tokens_img * num_heads * dim_head * raw_attn_output_split.scalar_size(),
batch_size, batch_size,
cudaMemcpyDeviceToDevice, hipMemcpyDeviceToDevice,
stream)); stream));
} }
...@@ -1112,8 +1112,8 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra ...@@ -1112,8 +1112,8 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra
kernels::mul_add_batch(attn_output, gate_msa, true, 0.0, hidden_states, true); kernels::mul_add_batch(attn_output, gate_msa, true, 0.0, hidden_states, true);
hidden_states = std::move(attn_output); hidden_states = std::move(attn_output);
nvtxRangePop(); roctxRangePop();
nvtxRangePushA("MLP"); roctxRangePushA("MLP");
spdlog::debug("attn_output={}", hidden_states.shape.str()); spdlog::debug("attn_output={}", hidden_states.shape.str());
...@@ -1138,7 +1138,7 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra ...@@ -1138,7 +1138,7 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra
kernels::mul_add_batch(ff_output, gate_mlp, true, 0.0, hidden_states, true); kernels::mul_add_batch(ff_output, gate_mlp, true, 0.0, hidden_states, true);
hidden_states = std::move(ff_output); hidden_states = std::move(ff_output);
nvtxRangePop(); roctxRangePop();
spdlog::debug("ff_output={}", hidden_states.shape.str()); spdlog::debug("ff_output={}", hidden_states.shape.str());
} }
...@@ -1148,7 +1148,7 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra ...@@ -1148,7 +1148,7 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra
} }
{ {
nvtxRangePushA("o_proj_context"); roctxRangePushA("o_proj_context");
auto &&[_, gate_msa, shift_mlp, scale_mlp, gate_mlp] = norm1_context_output; auto &&[_, gate_msa, shift_mlp, scale_mlp, gate_mlp] = norm1_context_output;
...@@ -1160,7 +1160,7 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra ...@@ -1160,7 +1160,7 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra
raw_attn_output_split = Tensor::allocate({batch_size, num_tokens_txt, num_heads * dim_head}, raw_attn_output_split = Tensor::allocate({batch_size, num_tokens_txt, num_heads * dim_head},
raw_attn_output.scalar_type(), raw_attn_output.scalar_type(),
raw_attn_output.device()); raw_attn_output.device());
checkCUDA(cudaMemcpy2DAsync(raw_attn_output_split.data_ptr(), checkCUDA(hipMemcpy2DAsync(raw_attn_output_split.data_ptr(),
num_tokens_txt * num_heads * dim_head * raw_attn_output_split.scalar_size(), num_tokens_txt * num_heads * dim_head * raw_attn_output_split.scalar_size(),
raw_attn_output.data_ptr<char>() + num_tokens_img_pad * num_heads * dim_head * raw_attn_output.data_ptr<char>() + num_tokens_img_pad * num_heads * dim_head *
raw_attn_output_split.scalar_size(), raw_attn_output_split.scalar_size(),
...@@ -1168,7 +1168,7 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra ...@@ -1168,7 +1168,7 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra
raw_attn_output.scalar_size(), raw_attn_output.scalar_size(),
num_tokens_txt * num_heads * dim_head * raw_attn_output_split.scalar_size(), num_tokens_txt * num_heads * dim_head * raw_attn_output_split.scalar_size(),
batch_size, batch_size,
cudaMemcpyDeviceToDevice, hipMemcpyDeviceToDevice,
stream)); stream));
} }
...@@ -1185,8 +1185,8 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra ...@@ -1185,8 +1185,8 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra
kernels::mul_add_batch(attn_output, gate_msa, true, 0.0, encoder_hidden_states, true); kernels::mul_add_batch(attn_output, gate_msa, true, 0.0, encoder_hidden_states, true);
encoder_hidden_states = std::move(attn_output); encoder_hidden_states = std::move(attn_output);
nvtxRangePop(); roctxRangePop();
nvtxRangePushA("MLP"); roctxRangePushA("MLP");
spdlog::debug("attn_output={}", encoder_hidden_states.shape.str()); spdlog::debug("attn_output={}", encoder_hidden_states.shape.str());
...@@ -1213,12 +1213,12 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra ...@@ -1213,12 +1213,12 @@ std::tuple<Tensor, Tensor, Tensor> JointTransformerBlock::forward_ip_adapter_bra
kernels::mul_add_batch(ff_output, gate_mlp, true, 0.0, encoder_hidden_states, true); kernels::mul_add_batch(ff_output, gate_mlp, true, 0.0, encoder_hidden_states, true);
encoder_hidden_states = std::move(ff_output); encoder_hidden_states = std::move(ff_output);
nvtxRangePop(); roctxRangePop();
spdlog::debug("ff_output={}", encoder_hidden_states.shape.str()); spdlog::debug("ff_output={}", encoder_hidden_states.shape.str());
} }
nvtxRangePop(); roctxRangePop();
return {hidden_states, encoder_hidden_states, q_heads}; return {hidden_states, encoder_hidden_states, q_heads};
} }
......
src/Linear.cpp
View file @ 1a8114bf
...@@ -5,7 +5,7 @@ ...@@ -5,7 +5,7 @@
#include "kernels/awq/gemv_awq.h" #include "kernels/awq/gemv_awq.h"
#include "kernels/dwconv.h" #include "kernels/dwconv.h"
#include <nvtx3/nvToolsExt.h> #include <nvtx3/roctracer/roctx.h>
using namespace nunchaku; using namespace nunchaku;
...@@ -117,7 +117,7 @@ GEMM_W4A4::GEMM_W4A4( ...@@ -117,7 +117,7 @@ GEMM_W4A4::GEMM_W4A4(
wtscale, "wtscale", ParamFlags::Optional)(wcscales, "wcscales", ParamFlags::Optional); wtscale, "wtscale", ParamFlags::Optional)(wcscales, "wcscales", ParamFlags::Optional);
#if NO_LORA_FUSION #if NO_LORA_FUSION
checkCUBLAS(cublasCreate(&handle)); checkCUBLAS(hipblasCreate(&handle));
#endif #endif
} }
...@@ -140,7 +140,7 @@ void GEMM_W4A4::loadParam(std::string key, Tensor &dst, Tensor src) { ...@@ -140,7 +140,7 @@ void GEMM_W4A4::loadParam(std::string key, Tensor &dst, Tensor src) {
} else if (key == "wtscale") { } else if (key == "wtscale") {
assert(src.numel() == 1); assert(src.numel() == 1);
if (src.dtype() == Tensor::BF16) { if (src.dtype() == Tensor::BF16) {
*dst.data_ptr<float>() = float(*src.data_ptr<__nv_bfloat16>()); *dst.data_ptr<float>() = float(*src.data_ptr<__hip_bfloat16>());
} else if (src.dtype() == Tensor::FP16) { } else if (src.dtype() == Tensor::FP16) {
*dst.data_ptr<float>() = float(*src.data_ptr<half>()); *dst.data_ptr<float>() = float(*src.data_ptr<half>());
} else if (src.dtype() == Tensor::FP32) { } else if (src.dtype() == Tensor::FP32) {
...@@ -242,15 +242,15 @@ void GEMM_W4A4::forward(Tensor x, ...@@ -242,15 +242,15 @@ void GEMM_W4A4::forward(Tensor x,
qact.is_unsigned, qact.is_unsigned,
this->lora_scales); this->lora_scales);
nvtxRangePushA("LoraUp"); roctxRangePushA("LoraUp");
static const half one = 1.0; static const half one = 1.0;
static const half zero = 0.0; static const half zero = 0.0;
// lora_up: [M, R] * [OC, R] => [M, OC] // lora_up: [M, R] * [OC, R] => [M, OC]
// cublas view: [OC, R] * [M, R]^T // cublas view: [OC, R] * [M, R]^T
checkCUBLAS(cublasHgemm(handle, checkCUBLAS(hipblasHgemm(handle,
CUBLAS_OP_T, HIPBLAS_OP_T,
CUBLAS_OP_N, HIPBLAS_OP_N,
this->out_features, this->out_features,
M, M,
this->lora_rank, this->lora_rank,
...@@ -263,7 +263,7 @@ void GEMM_W4A4::forward(Tensor x, ...@@ -263,7 +263,7 @@ void GEMM_W4A4::forward(Tensor x,
out.data_ptr<half>(), out.data_ptr<half>(),
this->out_features)); this->out_features));
nvtxRangePop(); roctxRangePop();
#endif #endif
} }
...@@ -380,7 +380,7 @@ GEMM_W4A4::forward_quant(QuantizedActivation qact, FuseOptions fuse, GEMM_W4A4 * ...@@ -380,7 +380,7 @@ GEMM_W4A4::forward_quant(QuantizedActivation qact, FuseOptions fuse, GEMM_W4A4 *
qact.is_unsigned, qact.is_unsigned,
this->lora_scales); this->lora_scales);
nvtxRangePushA("LoraUp"); roctxRangePushA("LoraUp");
static const half one = 1.0; static const half one = 1.0;
static const half zero = 0.0; static const half zero = 0.0;
...@@ -388,9 +388,9 @@ GEMM_W4A4::forward_quant(QuantizedActivation qact, FuseOptions fuse, GEMM_W4A4 * ...@@ -388,9 +388,9 @@ GEMM_W4A4::forward_quant(QuantizedActivation qact, FuseOptions fuse, GEMM_W4A4 *
// lora_up: [M, R] * [OC, R]^T => [M, OC] // lora_up: [M, R] * [OC, R]^T => [M, OC]
// cublas view: [R, OC]^T * [R, M] => [OC, M] // cublas view: [R, OC]^T * [R, M] => [OC, M]
// lora_up layout wrong? // lora_up layout wrong?
checkCUBLAS(cublasHgemm(handle, checkCUBLAS(hipblasHgemm(handle,
CUBLAS_OP_T, HIPBLAS_OP_T,
CUBLAS_OP_N, HIPBLAS_OP_N,
this->out_features, this->out_features,
M, M,
this->lora_rank, this->lora_rank,
...@@ -403,16 +403,16 @@ GEMM_W4A4::forward_quant(QuantizedActivation qact, FuseOptions fuse, GEMM_W4A4 * ...@@ -403,16 +403,16 @@ GEMM_W4A4::forward_quant(QuantizedActivation qact, FuseOptions fuse, GEMM_W4A4 *
out.data_ptr<half>(), out.data_ptr<half>(),
this->out_features)); this->out_features));
nvtxRangePop(); roctxRangePop();
if (fuse == FuseOptions::GELU_QUANT) { if (fuse == FuseOptions::GELU_QUANT) {
nvtxRangePushA("LoraDown"); roctxRangePushA("LoraDown");
// IC is for next lora (OC of this layer) // IC is for next lora (OC of this layer)
// lora_down: [M, IC] * [IC, R] => [M, R] // lora_down: [M, IC] * [IC, R] => [M, R]
// cublas view: [R, IC] * [IC, M] => [R, M] // cublas view: [R, IC] * [IC, M] => [R, M]
checkCUBLAS(cublasHgemm(handle, checkCUBLAS(hipblasHgemm(handle,
CUBLAS_OP_N, HIPBLAS_OP_N,
CUBLAS_OP_N, HIPBLAS_OP_N,
this->lora_rank, this->lora_rank,
M, M,
this->out_features, this->out_features,
...@@ -427,7 +427,7 @@ GEMM_W4A4::forward_quant(QuantizedActivation qact, FuseOptions fuse, GEMM_W4A4 * ...@@ -427,7 +427,7 @@ GEMM_W4A4::forward_quant(QuantizedActivation qact, FuseOptions fuse, GEMM_W4A4 *
out = {}; out = {};
nvtxRangePop(); roctxRangePop();
} }
#endif #endif
...@@ -473,13 +473,13 @@ GEMM_W4A4::QuantizedActivation GEMM_W4A4::quantize(Tensor x, bool fuse_glu) { ...@@ -473,13 +473,13 @@ GEMM_W4A4::QuantizedActivation GEMM_W4A4::quantize(Tensor x, bool fuse_glu) {
static const half one = 1.0; static const half one = 1.0;
static const half zero = 0.0; static const half zero = 0.0;
nvtxRangePushA("LoraDown"); roctxRangePushA("LoraDown");
// lora_down: [M, IC] * [IC, R] => [M, R] // lora_down: [M, IC] * [IC, R] => [M, R]
// cublas view: [R, IC] * [IC, M] // cublas view: [R, IC] * [IC, M]
checkCUBLAS(cublasHgemm(handle, checkCUBLAS(hipblasHgemm(handle,
CUBLAS_OP_N, HIPBLAS_OP_N,
CUBLAS_OP_N, HIPBLAS_OP_N,
this->lora_rank, this->lora_rank,
M, M,
this->in_features, this->in_features,
...@@ -492,7 +492,7 @@ GEMM_W4A4::QuantizedActivation GEMM_W4A4::quantize(Tensor x, bool fuse_glu) { ...@@ -492,7 +492,7 @@ GEMM_W4A4::QuantizedActivation GEMM_W4A4::quantize(Tensor x, bool fuse_glu) {
qact.lora_act.data_ptr<half>(), qact.lora_act.data_ptr<half>(),
this->lora_rank)); this->lora_rank));
nvtxRangePop(); roctxRangePop();
kernels::quantize_w4a4_act(x, qact.act, qact.ascales); kernels::quantize_w4a4_act(x, qact.act, qact.ascales);
......
src/Linear.h
View file @ 1a8114bf
...@@ -116,7 +116,7 @@ public: ...@@ -116,7 +116,7 @@ public:
Tensor wtscale; Tensor wtscale;
Tensor wcscales; Tensor wcscales;
cublasHandle_t handle; hipblasHandle_t handle;
}; };
class GEMM_W8A8 : public Module { class GEMM_W8A8 : public Module {
......
src/Module.h
View file @ 1a8114bf
...@@ -258,7 +258,7 @@ private: ...@@ -258,7 +258,7 @@ private:
waitEvent(eventLoadDone.get()); waitEvent(eventLoadDone.get());
funcCompute(layer); funcCompute(layer);
nextComputeDone = std::make_unique<CUDAEventWrapper>(); nextComputeDone = std::make_unique<CUDAEventWrapper>();
checkCUDA(cudaEventRecord(nextComputeDone->event, getCurrentCUDAStream())); checkCUDA(hipEventRecord(nextComputeDone->event, getCurrentHIPStreamMasqueradingAsCUDA()));
workaroundFlush(); workaroundFlush();
} }
...@@ -272,7 +272,7 @@ private: ...@@ -272,7 +272,7 @@ private:
funcLoad(layer + 1); funcLoad(layer + 1);
} }
nextLoadDone = std::make_unique<CUDAEventWrapper>(); nextLoadDone = std::make_unique<CUDAEventWrapper>();
checkCUDA(cudaEventRecord(nextLoadDone->event, getCurrentCUDAStream())); checkCUDA(hipEventRecord(nextLoadDone->event, getCurrentHIPStreamMasqueradingAsCUDA()));
workaroundFlush(); workaroundFlush();
} }
...@@ -287,7 +287,7 @@ private: ...@@ -287,7 +287,7 @@ private:
if (!event) { if (!event) {
return; return;
} }
checkCUDA(cudaStreamWaitEvent(getCurrentCUDAStream(), event->event)); checkCUDA(hipStreamWaitEvent(getCurrentHIPStreamMasqueradingAsCUDA(), event->event));
} }
// WDDM prevents multiple streams run concurrently // WDDM prevents multiple streams run concurrently
...@@ -312,12 +312,12 @@ private: ...@@ -312,12 +312,12 @@ private:
if (!needWorkaround) { if (!needWorkaround) {
return; return;
} }
cudaStreamQuery(getCurrentCUDAStream()); hipStreamQuery(getCurrentHIPStreamMasqueradingAsCUDA());
} }
void workaroundSynchronize() { void workaroundSynchronize() {
if (!needWorkaround) { if (!needWorkaround) {
return; return;
} }
checkCUDA(cudaEventSynchronize(eventComputeDone->event)); checkCUDA(hipEventSynchronize(eventComputeDone->event));
} }
}; };
src/SanaModel.cpp
View file @ 1a8114bf
...@@ -5,7 +5,7 @@ ...@@ -5,7 +5,7 @@
#include "flash_api.h" #include "flash_api.h"
#include "kernels/misc_kernels.h" #include "kernels/misc_kernels.h"
#include <nvtx3/nvToolsExt.h> #include <nvtx3/roctracer/roctx.h>
using spdlog::fmt_lib::format; using spdlog::fmt_lib::format;
using namespace nunchaku; using namespace nunchaku;
...@@ -241,9 +241,9 @@ Tensor SanaLinearTransformerBlock::forward(Tensor hidden_states, ...@@ -241,9 +241,9 @@ Tensor SanaLinearTransformerBlock::forward(Tensor hidden_states,
bool pag, bool pag,
bool cfg) { bool cfg) {
nvtxRangePushA("SanaLinearTransformerBlock"); roctxRangePushA("SanaLinearTransformerBlock");
nvtxRangePushA("chunk"); roctxRangePushA("chunk");
// Tensor ones = Tensor::ones({hidden_size}, Tensor::FP16, x.device()); // Tensor ones = Tensor::ones({hidden_size}, Tensor::FP16, x.device());
...@@ -262,10 +262,10 @@ Tensor SanaLinearTransformerBlock::forward(Tensor hidden_states, ...@@ -262,10 +262,10 @@ Tensor SanaLinearTransformerBlock::forward(Tensor hidden_states,
auto &&[shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp] = chunked; auto &&[shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp] = chunked;
// auto &&[shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp] = kernels::split_mod<6>(timestep); // auto &&[shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp] = kernels::split_mod<6>(timestep);
nvtxRangePop(); roctxRangePop();
{ {
nvtxRangePushA("LinearAttention"); roctxRangePushA("LinearAttention");
Tensor residual = hidden_states; Tensor residual = hidden_states;
Tensor norm_hidden_states = norm1.forward(hidden_states); Tensor norm_hidden_states = norm1.forward(hidden_states);
...@@ -279,11 +279,11 @@ Tensor SanaLinearTransformerBlock::forward(Tensor hidden_states, ...@@ -279,11 +279,11 @@ Tensor SanaLinearTransformerBlock::forward(Tensor hidden_states,
hidden_states = attn_output; hidden_states = attn_output;
nvtxRangePop(); roctxRangePop();
} }
{ {
nvtxRangePushA("CrossAttention"); roctxRangePushA("CrossAttention");
debug("norm_hidden_states_cross", hidden_states); debug("norm_hidden_states_cross", hidden_states);
Tensor attn_output = cross_attn.forward(hidden_states, encoder_hidden_states, cu_seqlens_img, cu_seqlens_txt); Tensor attn_output = cross_attn.forward(hidden_states, encoder_hidden_states, cu_seqlens_img, cu_seqlens_txt);
...@@ -293,11 +293,11 @@ Tensor SanaLinearTransformerBlock::forward(Tensor hidden_states, ...@@ -293,11 +293,11 @@ Tensor SanaLinearTransformerBlock::forward(Tensor hidden_states,
hidden_states = attn_output; hidden_states = attn_output;
nvtxRangePop(); roctxRangePop();
} }
{ {
nvtxRangePushA("Feed-forward"); roctxRangePushA("Feed-forward");
debug("hidden_states_ff", hidden_states); debug("hidden_states_ff", hidden_states);
Tensor norm_hidden_states = norm2.forward(hidden_states); Tensor norm_hidden_states = norm2.forward(hidden_states);
...@@ -311,10 +311,10 @@ Tensor SanaLinearTransformerBlock::forward(Tensor hidden_states, ...@@ -311,10 +311,10 @@ Tensor SanaLinearTransformerBlock::forward(Tensor hidden_states,
hidden_states = ff_output; hidden_states = ff_output;
nvtxRangePop(); roctxRangePop();
} }
nvtxRangePop(); roctxRangePop();
debug("hidden_states_out", hidden_states); debug("hidden_states_out", hidden_states);
......
src/Serialization.cpp
View file @ 1a8114bf
...@@ -121,15 +121,15 @@ SafeTensors::SafeTensors(const std::string &filename) { ...@@ -121,15 +121,15 @@ SafeTensors::SafeTensors(const std::string &filename) {
auto methodPrivate = [&]() { auto methodPrivate = [&]() {
this->mapped = std::make_unique<MMapImplPrivate>(filename); this->mapped = std::make_unique<MMapImplPrivate>(filename);
checkCUDA( checkCUDA(
cudaHostRegister(const_cast<char *>(this->mapped->data()), this->mapped->size(), cudaHostRegisterPortable)); hipHostRegister(const_cast<char *>(this->mapped->data()), this->mapped->size(), hipHostRegisterPortable));
this->hostRegistered = true; this->hostRegistered = true;
this->memoryPinned = true; this->memoryPinned = true;
}; };
auto methodMio = [&]() { auto methodMio = [&]() {
this->mapped = std::make_unique<MMapImplMio>(filename); this->mapped = std::make_unique<MMapImplMio>(filename);
checkCUDA(cudaHostRegister(const_cast<char *>(this->mapped->data()), checkCUDA(hipHostRegister(const_cast<char *>(this->mapped->data()),
this->mapped->size(), this->mapped->size(),
cudaHostRegisterPortable | cudaHostRegisterReadOnly)); hipHostRegisterPortable | hipHostRegisterReadOnly));
this->hostRegistered = true; this->hostRegistered = true;
this->memoryPinned = true; this->memoryPinned = true;
}; };
...@@ -183,8 +183,8 @@ SafeTensors::SafeTensors(const std::string &filename) { ...@@ -183,8 +183,8 @@ SafeTensors::SafeTensors(const std::string &filename) {
SafeTensors::~SafeTensors() { SafeTensors::~SafeTensors() {
if (this->hostRegistered) { if (this->hostRegistered) {
if (cudaHostUnregister(const_cast<char *>(this->mapped->data())) != cudaSuccess) { if (hipHostUnregister(const_cast<char *>(this->mapped->data())) != hipSuccess) {
spdlog::warn("cudaHostUnregister failed: {}", cudaGetErrorString(cudaGetLastError())); spdlog::warn("hipHostUnregister failed: {}", hipGetErrorString(hipGetLastError()));
} }
} }
} }
......
src/Serialization.h
View file @ 1a8114bf
...@@ -9,17 +9,17 @@ public: ...@@ -9,17 +9,17 @@ public:
this->size = size; this->size = size;
this->device.type = Device::CPU; this->device.type = Device::CPU;
this->ptr = ptr; this->ptr = ptr;
// auto ret = cudaHostRegister(ptr, size, cudaHostRegisterPortable | cudaHostRegisterReadOnly); // auto ret = hipHostRegister(ptr, size, hipHostRegisterPortable | hipHostRegisterReadOnly);
// if (ret == cudaSuccess) { // if (ret == hipSuccess) {
// this->registered = true; // this->registered = true;
// } else { // } else {
// log(std::format("cudaHostRegister failed at {:p} (size={}): {}", ptr, size, // log(std::format("hipHostRegister failed at {:p} (size={}): {}", ptr, size,
// cudaGetErrorString(cudaGetLastError()))); this->registered = false; // hipGetErrorString(hipGetLastError()))); this->registered = false;
// } // }
} }
virtual ~BufferMMap() { virtual ~BufferMMap() {
// if (registered) { // if (registered) {
// checkCUDA(cudaHostUnregister(ptr)); // checkCUDA(hipHostUnregister(ptr));
// } // }
} }
......
src/Tensor.h
View file @ 1a8114bf
#include "hip/hip_runtime.h"
#pragma once #pragma once
#include "common.h" #include "common.h"
...@@ -74,10 +75,10 @@ public: ...@@ -74,10 +75,10 @@ public:
BufferHost(size_t size) { BufferHost(size_t size) {
this->size = size; this->size = size;
this->device.type = Device::CPU; this->device.type = Device::CPU;
checkCUDA(cudaHostAlloc(&this->ptr, size, cudaHostAllocPortable)); checkCUDA(hipHostMalloc(&this->ptr, size, hipHostMallocPortable));
} }
virtual ~BufferHost() { virtual ~BufferHost() {
checkCUDA(cudaFreeHost(this->ptr)); checkCUDA(hipHostFree(this->ptr));
} }
}; };
...@@ -86,20 +87,20 @@ public: ...@@ -86,20 +87,20 @@ public:
BufferCUDA(size_t size) { BufferCUDA(size_t size) {
this->size = size; this->size = size;
this->device.type = Device::CUDA; this->device.type = Device::CUDA;
// checkCUDA(cudaGetDevice(&this->device.idx)); // checkCUDA(hipGetDevice(&this->device.idx));
this->device.idx = CUDADeviceContext::getDevice(); this->device.idx = CUDADeviceContext::getDevice();
if (size == 0) { if (size == 0) {
this->ptr = nullptr; this->ptr = nullptr;
} }
// TODO: buffer used in multiple streams? // TODO: buffer used in multiple streams?
checkCUDA(cudaMallocAsync(&this->ptr, size, getCurrentCUDAStream())); checkCUDA(hipMallocAsync(&this->ptr, size, getCurrentHIPStreamMasqueradingAsCUDA()));
} }
virtual ~BufferCUDA() { virtual ~BufferCUDA() {
if (this->size == 0) { if (this->size == 0) {
assert(!this->ptr); assert(!this->ptr);
return; return;
} }
checkCUDA(cudaFreeAsync(this->ptr, getCurrentCUDAStream())); checkCUDA(hipFreeAsync(this->ptr, getCurrentHIPStreamMasqueradingAsCUDA()));
} }
virtual bool isAsyncBuffer() override { virtual bool isAsyncBuffer() override {
return true; return true;
...@@ -111,11 +112,11 @@ public: ...@@ -111,11 +112,11 @@ public:
BufferCUDASync(size_t size) { BufferCUDASync(size_t size) {
this->size = size; this->size = size;
this->device.type = Device::CUDA; this->device.type = Device::CUDA;
checkCUDA(cudaGetDevice(&this->device.idx)); checkCUDA(hipGetDevice(&this->device.idx));
checkCUDA(cudaMalloc(&this->ptr, size)); checkCUDA(hipMalloc(&this->ptr, size));
} }
virtual ~BufferCUDASync() { virtual ~BufferCUDASync() {
checkCUDA(cudaFree(this->ptr)); checkCUDA(hipFree(this->ptr));
} }
}; };
...@@ -415,8 +416,8 @@ public: ...@@ -415,8 +416,8 @@ public:
Tensor &zero_() { Tensor &zero_() {
assert(this->is_contiguous()); assert(this->is_contiguous());
checkCUDA(cudaMemsetAsync( checkCUDA(hipMemsetAsync(
data_ptr<char>() + shape.offset * scalar_size(), 0, shape.size() * scalar_size(), getCurrentCUDAStream())); data_ptr<char>() + shape.offset * scalar_size(), 0, shape.size() * scalar_size(), getCurrentHIPStreamMasqueradingAsCUDA()));
return *this; return *this;
} }
Tensor &copy_(Tensor other) { Tensor &copy_(Tensor other) {
...@@ -444,13 +445,13 @@ public: ...@@ -444,13 +445,13 @@ public:
return *this; return *this;
} }
lockBuffer(this->buffer, getCurrentCUDAStream()); lockBuffer(this->buffer, getCurrentHIPStreamMasqueradingAsCUDA());
lockBuffer(other.buffer, getCurrentCUDAStream()); lockBuffer(other.buffer, getCurrentHIPStreamMasqueradingAsCUDA());
checkCUDA(cudaMemcpyAsync(data_ptr<char>(), checkCUDA(hipMemcpyAsync(data_ptr<char>(),
other.data_ptr<char>(), other.data_ptr<char>(),
shape.size() * scalar_size(), shape.size() * scalar_size(),
getCopyKind(this->device(), other.device()), getCopyKind(this->device(), other.device()),
getCurrentCUDAStream())); getCurrentHIPStreamMasqueradingAsCUDA()));
return *this; return *this;
} }
...@@ -487,7 +488,7 @@ public: ...@@ -487,7 +488,7 @@ public:
} else if (device.type == Device::CUDA) { } else if (device.type == Device::CUDA) {
CUDADeviceContext ctx(device.idx); CUDADeviceContext ctx(device.idx);
checkCUDA( checkCUDA(
cudaMemsetAsync(result.buffer->getPtr(), 0xCC, result.buffer->getSize(), getCurrentCUDAStream())); hipMemsetAsync(result.buffer->getPtr(), 0xCC, result.buffer->getSize(), getCurrentHIPStreamMasqueradingAsCUDA()));
} }
} }
...@@ -502,7 +503,7 @@ public: ...@@ -502,7 +503,7 @@ public:
static Tensor ones(TensorShape shape, ScalarType scalarType, Device device) { static Tensor ones(TensorShape shape, ScalarType scalarType, Device device) {
Tensor result = allocate(shape, scalarType, device); Tensor result = allocate(shape, scalarType, device);
// FIXME FIXME FIXME // FIXME FIXME FIXME
checkCUDA(cudaMemsetAsync(result.buffer->getPtr(), 1, result.buffer->getSize(), getCurrentCUDAStream())); checkCUDA(hipMemsetAsync(result.buffer->getPtr(), 1, result.buffer->getSize(), getCurrentHIPStreamMasqueradingAsCUDA()));
return result; return result;
} }
static Tensor static Tensor
...@@ -522,18 +523,18 @@ public: ...@@ -522,18 +523,18 @@ public:
Tensor result = allocate(this->shape.dataExtent, this->scalarType, device); Tensor result = allocate(this->shape.dataExtent, this->scalarType, device);
result.copy_(*this); result.copy_(*this);
// lockBuffer(this->buffer, getCurrentCUDAStream()); // lockBuffer(this->buffer, getCurrentHIPStreamMasqueradingAsCUDA());
// lockBuffer(result.buffer, getCurrentCUDAStream()); // lockBuffer(result.buffer, getCurrentHIPStreamMasqueradingAsCUDA());
// checkCUDA(cudaMemcpyAsync(result.data_ptr(), this->data_ptr(), result.buffer->getSize(), cudaMemcpyDefault, // checkCUDA(hipMemcpyAsync(result.data_ptr(), this->data_ptr(), result.buffer->getSize(), hipMemcpyDefault,
// getCurrentCUDAStream())); if (this->device().type == Device::CPU && device.type == Device::CUDA) { // getCurrentHIPStreamMasqueradingAsCUDA())); if (this->device().type == Device::CPU && device.type == Device::CUDA) {
// checkCUDA(cudaMemcpyAsync(result.data_ptr(), this->data_ptr(), result.buffer->getSize(), // checkCUDA(hipMemcpyAsync(result.data_ptr(), this->data_ptr(), result.buffer->getSize(),
// cudaMemcpyHostToDevice, getCurrentCUDAStream())); // hipMemcpyHostToDevice, getCurrentHIPStreamMasqueradingAsCUDA()));
// } else if (this->device().type == Device::CUDA && device.type == Device::CPU) { // } else if (this->device().type == Device::CUDA && device.type == Device::CPU) {
// checkCUDA(cudaMemcpyAsync(result.data_ptr(), this->data_ptr(), result.buffer->getSize(), // checkCUDA(hipMemcpyAsync(result.data_ptr(), this->data_ptr(), result.buffer->getSize(),
// cudaMemcpyDeviceToHost, getCurrentCUDAStream())); // hipMemcpyDeviceToHost, getCurrentHIPStreamMasqueradingAsCUDA()));
// } else { // } else {
// checkCUDA(cudaMemcpyAsync(result.data_ptr(), this->data_ptr(), result.buffer->getSize(), // checkCUDA(hipMemcpyAsync(result.data_ptr(), this->data_ptr(), result.buffer->getSize(),
// cudaMemcpyDefault, getCurrentCUDAStream())); // hipMemcpyDefault, getCurrentHIPStreamMasqueradingAsCUDA()));
// } // }
return result; return result;
} }
...@@ -548,38 +549,38 @@ public: ...@@ -548,38 +549,38 @@ public:
// auto shapeOut = this->shape; // auto shapeOut = this->shape;
// shapeOut[dim] = upper_bound - lower_bound; // shapeOut[dim] = upper_bound - lower_bound;
// assert(dst.shape.data == shapeOut.data); // assert(dst.shape.data == shapeOut.data);
// checkCUDA(cudaMemcpy2DAsync( // checkCUDA(hipMemcpy2DAsync(
// dst. // dst.
// )); // ));
// } // }
private: private:
static cudaMemcpyKind getCopyKind(Device dst, Device src) { static hipMemcpyKind getCopyKind(Device dst, Device src) {
if (src.type == Device::CPU && dst.type == Device::CUDA) { if (src.type == Device::CPU && dst.type == Device::CUDA) {
return cudaMemcpyHostToDevice; return hipMemcpyHostToDevice;
} }
if (src.type == Device::CUDA && dst.type == Device::CPU) { if (src.type == Device::CUDA && dst.type == Device::CPU) {
return cudaMemcpyDeviceToHost; return hipMemcpyDeviceToHost;
} }
if (src.type == Device::CUDA && dst.type == Device::CUDA) { if (src.type == Device::CUDA && dst.type == Device::CUDA) {
return cudaMemcpyDeviceToDevice; return hipMemcpyDeviceToDevice;
} }
if (src.type == Device::CPU && dst.type == Device::CPU) { if (src.type == Device::CPU && dst.type == Device::CPU) {
return cudaMemcpyHostToHost; return hipMemcpyHostToHost;
} }
return cudaMemcpyDefault; return hipMemcpyDefault;
} }
// static bool isAsyncBuffer(Buffer *buffer) { // static bool isAsyncBuffer(Buffer *buffer) {
// return dynamic_cast<BufferCUDA *>(buffer); // return dynamic_cast<BufferCUDA *>(buffer);
// } // }
static inline std::map<cudaStream_t, std::set<std::shared_ptr<Buffer>>> lockedBuffers; static inline std::map<hipStream_t, std::set<std::shared_ptr<Buffer>>> lockedBuffers;
public: public:
// before launching an async operation, make sure to lock the buffer in case the buffer is freed before GPU // before launching an async operation, make sure to lock the buffer in case the buffer is freed before GPU
// completes // completes
static void lockBuffer(std::shared_ptr<Buffer> buffer, cudaStream_t stream) { static void lockBuffer(std::shared_ptr<Buffer> buffer, hipStream_t stream) {
if (!buffer->isAsyncBuffer()) { if (!buffer->isAsyncBuffer()) {
lockedBuffers[stream].insert(buffer); lockedBuffers[stream].insert(buffer);
} }
...@@ -589,16 +590,16 @@ public: ...@@ -589,16 +590,16 @@ public:
static void unlockBuffers() { static void unlockBuffers() {
lockedBuffers.clear(); lockedBuffers.clear();
} }
static void unlockBuffers(cudaStream_t stream) { static void unlockBuffers(hipStream_t stream) {
lockedBuffers[stream].clear(); lockedBuffers[stream].clear();
} }
static void synchronizeDevice() { static void synchronizeDevice() {
checkCUDA(cudaDeviceSynchronize()); checkCUDA(hipDeviceSynchronize());
unlockBuffers(); unlockBuffers();
} }
static void synchronizeStream(cudaStream_t stream) { static void synchronizeStream(hipStream_t stream) {
checkCUDA(cudaStreamSynchronize(stream)); checkCUDA(hipStreamSynchronize(stream));
unlockBuffers(stream); unlockBuffers(stream);
} }
}; };
......
src/common.h
View file @ 1a8114bf
...@@ -19,47 +19,47 @@ ...@@ -19,47 +19,47 @@
#include <optional> #include <optional>
#include <chrono> #include <chrono>
#include <functional> #include <functional>
#include <cuda_runtime_api.h> #include <hip/hip_runtime_api.h>
#include <cublas_v2.h> #include <hipblas/hipblas.h>
#include <spdlog/spdlog.h> #include <spdlog/spdlog.h>
class CUDAError : public std::runtime_error { class CUDAError : public std::runtime_error {
public: public:
CUDAError(cudaError_t errorCode, std::source_location location) CUDAError(hipError_t errorCode, std::source_location location)
: std::runtime_error(format(errorCode, location)), errorCode(errorCode), location(location) {} : std::runtime_error(format(errorCode, location)), errorCode(errorCode), location(location) {}
public: public:
const cudaError_t errorCode; const hipError_t errorCode;
const std::source_location location; const std::source_location location;
private: private:
static std::string format(cudaError_t errorCode, std::source_location location) { static std::string format(hipError_t errorCode, std::source_location location) {
return spdlog::fmt_lib::format( return spdlog::fmt_lib::format(
"CUDA error: {} (at {}:{})", cudaGetErrorString(errorCode), location.file_name(), location.line()); "CUDA error: {} (at {}:{})", hipGetErrorString(errorCode), location.file_name(), location.line());
} }
}; };
inline cudaError_t checkCUDA(cudaError_t retValue, inline hipError_t checkCUDA(hipError_t retValue,
const std::source_location location = std::source_location::current()) { const std::source_location location = std::source_location::current()) {
if (retValue != cudaSuccess) { if (retValue != hipSuccess) {
(void)cudaGetLastError(); (void)hipGetLastError();
throw CUDAError(retValue, location); throw CUDAError(retValue, location);
} }
return retValue; return retValue;
} }
inline cublasStatus_t checkCUBLAS(cublasStatus_t retValue, inline hipblasStatus_t checkCUBLAS(hipblasStatus_t retValue,
const std::source_location location = std::source_location::current()) { const std::source_location location = std::source_location::current()) {
if (retValue != CUBLAS_STATUS_SUCCESS) { if (retValue != HIPBLAS_STATUS_SUCCESS) {
throw std::runtime_error(spdlog::fmt_lib::format( throw std::runtime_error(spdlog::fmt_lib::format(
"CUBLAS error: {} (at {}:{})", cublasGetStatusString(retValue), location.file_name(), location.line())); "CUBLAS error: {} (at {}:{})", rocblas_status_to_string(retValue), location.file_name(), location.line()));
} }
return retValue; return retValue;
} }
inline thread_local std::stack<cudaStream_t> stackCUDAStreams; inline thread_local std::stack<hipStream_t> stackCUDAStreams;
inline cudaStream_t getCurrentCUDAStream() { inline hipStream_t getCurrentHIPStreamMasqueradingAsCUDA() {
if (stackCUDAStreams.empty()) { if (stackCUDAStreams.empty()) {
return 0; return 0;
} }
...@@ -67,9 +67,9 @@ inline cudaStream_t getCurrentCUDAStream() { ...@@ -67,9 +67,9 @@ inline cudaStream_t getCurrentCUDAStream() {
} }
struct CUDAStreamContext { struct CUDAStreamContext {
cudaStream_t stream; hipStream_t stream;
CUDAStreamContext(cudaStream_t stream) : stream(stream) { CUDAStreamContext(hipStream_t stream) : stream(stream) {
stackCUDAStreams.push(stream); stackCUDAStreams.push(stream);
} }
CUDAStreamContext(const CUDAStreamContext &) = delete; CUDAStreamContext(const CUDAStreamContext &) = delete;
...@@ -82,30 +82,30 @@ struct CUDAStreamContext { ...@@ -82,30 +82,30 @@ struct CUDAStreamContext {
}; };
struct CUDAStreamWrapper { struct CUDAStreamWrapper {
cudaStream_t stream; hipStream_t stream;
CUDAStreamWrapper() { CUDAStreamWrapper() {
checkCUDA(cudaStreamCreate(&stream)); checkCUDA(hipStreamCreate(&stream));
} }
CUDAStreamWrapper(const CUDAStreamWrapper &) = delete; CUDAStreamWrapper(const CUDAStreamWrapper &) = delete;
CUDAStreamWrapper(CUDAStreamWrapper &&) = delete; CUDAStreamWrapper(CUDAStreamWrapper &&) = delete;
~CUDAStreamWrapper() { ~CUDAStreamWrapper() {
checkCUDA(cudaStreamDestroy(stream)); checkCUDA(hipStreamDestroy(stream));
} }
}; };
struct CUDAEventWrapper { struct CUDAEventWrapper {
cudaEvent_t event; hipEvent_t event;
CUDAEventWrapper(unsigned int flags = cudaEventDefault) { CUDAEventWrapper(unsigned int flags = hipEventDefault) {
checkCUDA(cudaEventCreateWithFlags(&event, flags)); checkCUDA(hipEventCreateWithFlags(&event, flags));
} }
CUDAEventWrapper(const CUDAEventWrapper &) = delete; CUDAEventWrapper(const CUDAEventWrapper &) = delete;
CUDAEventWrapper(CUDAEventWrapper &&) = delete; CUDAEventWrapper(CUDAEventWrapper &&) = delete;
~CUDAEventWrapper() { ~CUDAEventWrapper() {
checkCUDA(cudaEventDestroy(event)); checkCUDA(hipEventDestroy(event));
} }
}; };
...@@ -162,7 +162,7 @@ public: ...@@ -162,7 +162,7 @@ public:
static int getDevice() { static int getDevice() {
int idx = -1; int idx = -1;
if (cacheDisabled() || currentDeviceCache < 0) { if (cacheDisabled() || currentDeviceCache < 0) {
checkCUDA(cudaGetDevice(&idx)); checkCUDA(hipGetDevice(&idx));
} else { } else {
idx = currentDeviceCache; idx = currentDeviceCache;
} }
...@@ -177,7 +177,7 @@ private: ...@@ -177,7 +177,7 @@ private:
if (!cacheDisabled() && currentDeviceCache == idx) { if (!cacheDisabled() && currentDeviceCache == idx) {
return; return;
} }
checkCUDA(cudaSetDevice(idx)); checkCUDA(hipSetDevice(idx));
currentDeviceCache = cacheDisabled() ? -1 : idx; currentDeviceCache = cacheDisabled() ? -1 : idx;
} }
...@@ -190,13 +190,13 @@ private: ...@@ -190,13 +190,13 @@ private:
} }
}; };
inline cudaDeviceProp *getCurrentDeviceProperties() { inline hipDeviceProp_t *getCurrentDeviceProperties() {
static thread_local std::map<int, cudaDeviceProp> props; static thread_local std::map<int, hipDeviceProp_t> props;
int deviceId = CUDADeviceContext::getDevice(); int deviceId = CUDADeviceContext::getDevice();
if (!props.contains(deviceId)) { if (!props.contains(deviceId)) {
cudaDeviceProp prop; hipDeviceProp_t prop;
checkCUDA(cudaGetDeviceProperties(&prop, deviceId)); checkCUDA(hipGetDeviceProperties(&prop, deviceId));
props[deviceId] = prop; props[deviceId] = prop;
} }
return &props.at(deviceId); return &props.at(deviceId);
...@@ -217,16 +217,16 @@ constexpr int log2Up(T value) { ...@@ -217,16 +217,16 @@ constexpr int log2Up(T value) {
} }
struct CUBLASWrapper { struct CUBLASWrapper {
cublasHandle_t handle = nullptr; hipblasHandle_t handle = nullptr;
CUBLASWrapper() { CUBLASWrapper() {
checkCUBLAS(cublasCreate(&handle)); checkCUBLAS(hipblasCreate(&handle));
} }
CUBLASWrapper(CUBLASWrapper &&) = delete; CUBLASWrapper(CUBLASWrapper &&) = delete;
CUBLASWrapper(const CUBLASWrapper &&) = delete; CUBLASWrapper(const CUBLASWrapper &&) = delete;
~CUBLASWrapper() { ~CUBLASWrapper() {
if (handle) { if (handle) {
checkCUBLAS(cublasDestroy(handle)); checkCUBLAS(hipblasDestroy(handle));
} }
} }
}; };
......
src/interop/torch.cpp
View file @ 1a8114bf
#include "torch.h" #include "torch.h"
#include <ATen/cuda/CUDAContext.h> #include <ATen/hip/HIPContext.h>
using spdlog::fmt_lib::format; using spdlog::fmt_lib::format;
...@@ -37,7 +37,7 @@ Tensor from_torch(at::Tensor input) { ...@@ -37,7 +37,7 @@ Tensor from_torch(at::Tensor input) {
result.scalarType = mapType.at(input.scalar_type()); result.scalarType = mapType.at(input.scalar_type());
result.buffer = std::make_shared<BufferTorchTensor>(std::move(input)); result.buffer = std::make_shared<BufferTorchTensor>(std::move(input));
Tensor::lockBuffer(result.buffer, getCurrentCUDAStream()); Tensor::lockBuffer(result.buffer, getCurrentHIPStreamMasqueradingAsCUDA());
return result; return result;
} }
...@@ -76,10 +76,10 @@ at::Tensor to_torch(Tensor input) { ...@@ -76,10 +76,10 @@ at::Tensor to_torch(Tensor input) {
} }
TorchOpContext::TorchOpContext() { TorchOpContext::TorchOpContext() {
stackCUDAStreams.push(at::cuda::getCurrentCUDAStream().stream()); stackCUDAStreams.push(at::hip::getCurrentHIPStreamMasqueradingAsCUDA().stream());
} }
TorchOpContext::~TorchOpContext() { TorchOpContext::~TorchOpContext() {
assert(stackCUDAStreams.top() == at::cuda::getCurrentCUDAStream().stream()); assert(stackCUDAStreams.top() == at::hip::getCurrentHIPStreamMasqueradingAsCUDA().stream());
stackCUDAStreams.pop(); stackCUDAStreams.pop();
} }
src/kernels/activation_kernels.cu → src/kernels/activation_kernels.hip
View file @ 1a8114bf
#include "hip/hip_runtime.h"
#include "activation_kernels_impl.cuh" #include "activation_kernels_impl.cuh"
#include "activation_kernels.h" #include "activation_kernels.h"
#include "dispatch_utils.h" #include "dispatch_utils.h"
...@@ -8,10 +9,10 @@ ...@@ -8,10 +9,10 @@
int num_tokens = input.numel() / d; \ int num_tokens = input.numel() / d; \
dim3 grid(num_tokens); \ dim3 grid(num_tokens); \
dim3 block(std::min(d, 1024)); \ dim3 block(std::min(d, 1024)); \
const cudaStream_t stream = getCurrentCUDAStream(); \ const hipStream_t stream = getCurrentHIPStreamMasqueradingAsCUDA(); \
VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "activation_kernel", [&] { \ VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "activation_kernel", [&] { \
vllm::activation_kernel<scalar_t, KERNEL<scalar_t>> \ hipLaunchKernelGGL(( vllm::activation_kernel<scalar_t, KERNEL<scalar_t>>) \
<<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \ , dim3(grid), dim3(block), 0, stream, out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
}); });
void silu_and_mul(Tensor &out, // [..., d] void silu_and_mul(Tensor &out, // [..., d]
...@@ -21,14 +22,14 @@ void silu_and_mul(Tensor &out, // [..., d] ...@@ -21,14 +22,14 @@ void silu_and_mul(Tensor &out, // [..., d]
int d = input.size(-1) / 2; int d = input.size(-1) / 2;
dim3 grid(num_tokens); dim3 grid(num_tokens);
dim3 block(std::min(d, 1024)); dim3 block(std::min(d, 1024));
const cudaStream_t stream = getCurrentCUDAStream(); const hipStream_t stream = getCurrentHIPStreamMasqueradingAsCUDA();
// dispatchFloat(input.scalar_type(), [&]<typename scalar_t>() { // dispatchFloat(input.scalar_type(), [&]<typename scalar_t>() {
// vllm::silu_and_mul_kernel<scalar_t><<<grid, block, 0, stream>>>( // vllm::silu_and_mul_kernel<scalar_t><<<grid, block, 0, stream>>>(
// out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); // out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);
// }); // });
VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "silu_and_mul_kernel", [&] { VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "silu_and_mul_kernel", [&] {
vllm::silu_and_mul_kernel<scalar_t> hipLaunchKernelGGL(( vllm::silu_and_mul_kernel<scalar_t>)
<<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); , dim3(grid), dim3(block), 0, stream, out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);
}); });
} }
...@@ -41,8 +42,8 @@ void invoke_dequant_silu_and_mul_quant(Tensor &out, // [..., d] ...@@ -41,8 +42,8 @@ void invoke_dequant_silu_and_mul_quant(Tensor &out, // [..., d]
int d = input.size(-1) / 2; int d = input.size(-1) / 2;
dim3 grid(num_tokens); dim3 grid(num_tokens);
dim3 block(std::min(d, 1024)); dim3 block(std::min(d, 1024));
const cudaStream_t stream = getCurrentCUDAStream(); const hipStream_t stream = getCurrentHIPStreamMasqueradingAsCUDA();
vllm::dequant_silu_and_mul_quant_kernel<float, false><<<grid, block, 0, stream>>>( hipLaunchKernelGGL(( vllm::dequant_silu_and_mul_quant_kernel<float, false>), dim3(grid), dim3(block), 0, stream,
out.data_ptr<int8_t>(), input.data_ptr<int32_t>(), d, scale_gate, scale_up, scale_out); out.data_ptr<int8_t>(), input.data_ptr<int32_t>(), d, scale_gate, scale_up, scale_out);
} }
...@@ -57,8 +58,8 @@ void invoke_dequant_silu_and_mul_quant(Tensor &out, // [..., d] ...@@ -57,8 +58,8 @@ void invoke_dequant_silu_and_mul_quant(Tensor &out, // [..., d]
int d = input.size(-1) / 2; int d = input.size(-1) / 2;
dim3 grid(num_tokens); dim3 grid(num_tokens);
dim3 block(std::min(d, 1024)); dim3 block(std::min(d, 1024));
const cudaStream_t stream = getCurrentCUDAStream(); const hipStream_t stream = getCurrentHIPStreamMasqueradingAsCUDA();
vllm::dequant_silu_and_mul_quant_kernel<float *, true><<<grid, block, 0, stream>>>(out.data_ptr<int8_t>(), hipLaunchKernelGGL(( vllm::dequant_silu_and_mul_quant_kernel<float *, true>), dim3(grid), dim3(block), 0, stream, out.data_ptr<int8_t>(),
input.data_ptr<int32_t>(), input.data_ptr<int32_t>(),
d, d,
scale_gate, scale_gate,
......
src/kernels/activation_kernels_impl.cuh
View file @ 1a8114bf
#include "hip/hip_runtime.h"
#include "utils.cuh" #include "utils.cuh"
#include "reduction_utils.cuh" #include "reduction_utils.cuh"
......
src/kernels/awq/dequantize.cuh
View file @ 1a8114bf
...@@ -11,7 +11,7 @@ https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutl ...@@ -11,7 +11,7 @@ https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutl
*/ */
#pragma once #pragma once
#include <cuda_fp16.h> #include <hip/hip_fp16.h>
#include <cstdint> #include <cstdint>
__forceinline__ __device__ void dequantize_s4_to_fp16x2(half2 const &source, uint4 *result) { __forceinline__ __device__ void dequantize_s4_to_fp16x2(half2 const &source, uint4 *result) {
...@@ -75,14 +75,14 @@ __forceinline__ __device__ void dequantize_s4_to_fp16x2(half2 const &source, uin ...@@ -75,14 +75,14 @@ __forceinline__ __device__ void dequantize_s4_to_fp16x2(half2 const &source, uin
asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" : "=r"(h[3]) : "r"(h[3]), "r"(ONE_SIXTEENTH), "r"(NEG_64)); asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" : "=r"(h[3]) : "r"(h[3]), "r"(ONE_SIXTEENTH), "r"(NEG_64));
} }
__forceinline__ __device__ void dequantize_s4_to_fp16x2(__nv_bfloat162 const &source, uint4 *result) { __forceinline__ __device__ void dequantize_s4_to_fp16x2(__hip_bfloat162 const &source, uint4 *result) {
// dequantize_s4_to_fp16x2(reinterpret_cast<const half2 &>(source), result); // dequantize_s4_to_fp16x2(reinterpret_cast<const half2 &>(source), result);
// *reinterpret_cast<__nv_bfloat162 *>(&result->x) = cuda_cast<__nv_bfloat162>(*reinterpret_cast<half2 // *reinterpret_cast<__hip_bfloat162 *>(&result->x) = cuda_cast<__hip_bfloat162>(*reinterpret_cast<half2
// *>(&result->x)); *reinterpret_cast<__nv_bfloat162 *>(&result->y) = // *>(&result->x)); *reinterpret_cast<__hip_bfloat162 *>(&result->y) =
// cuda_cast<__nv_bfloat162>(*reinterpret_cast<half2 *>(&result->y)); *reinterpret_cast<__nv_bfloat162 // cuda_cast<__hip_bfloat162>(*reinterpret_cast<half2 *>(&result->y)); *reinterpret_cast<__hip_bfloat162
// *>(&result->z) = cuda_cast<__nv_bfloat162>(*reinterpret_cast<half2 *>(&result->z)); // *>(&result->z) = cuda_cast<__hip_bfloat162>(*reinterpret_cast<half2 *>(&result->z));
// *reinterpret_cast<__nv_bfloat162 *>(&result->w) = cuda_cast<__nv_bfloat162>(*reinterpret_cast<half2 // *reinterpret_cast<__hip_bfloat162 *>(&result->w) = cuda_cast<__hip_bfloat162>(*reinterpret_cast<half2
// *>(&result->w)); // *>(&result->w));
// return; // return;
......
src/kernels/awq/gemm_awq.cu → src/kernels/awq/gemm_awq.hip
View file @ 1a8114bf
#include <cuda_fp16.h> #include "hip/hip_runtime.h"
#include <cuda_bf16.h> #include <hip/hip_fp16.h>
#include <hip/hip_bf16.h>
#include "semaphore.h" #include "semaphore.h"
#include "gemm_awq.h" #include "gemm_awq.h"
// #include "../../../nunchaku/csrc/quantization/dequantize.cuh" // #include "../../../nunchaku/csrc/quantization/dequantize.cuh"
...@@ -46,8 +47,8 @@ ...@@ -46,8 +47,8 @@
dim3 num_blocks((num_out_feats + CTA_M - 1) / CTA_M * j_factors1 * SPLITK); \ dim3 num_blocks((num_out_feats + CTA_M - 1) / CTA_M * j_factors1 * SPLITK); \
dim3 threads_per_block(WARP_SIZE, NUM_WARPS); \ dim3 threads_per_block(WARP_SIZE, NUM_WARPS); \
auto kernel_func = gemm_w4a16_T1<f16_t, CTA_M, CTA_N, CTA_K, WARP_M, WARP_N, WARP_K, STAGES, G, SPLITK>; \ auto kernel_func = gemm_w4a16_T1<f16_t, CTA_M, CTA_N, CTA_K, WARP_M, WARP_N, WARP_K, STAGES, G, SPLITK>; \
cudaFuncSetAttribute(kernel_func, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemByteSize); \ hipFuncSetAttribute(kernel_func, hipFuncAttributeMaxDynamicSharedMemorySize, kSmemByteSize); \
kernel_func<<<num_blocks, threads_per_block, kSmemByteSize>>>( \ hipLaunchKernelGGL(( kernel_func), dim3(num_blocks), dim3(threads_per_block), kSmemByteSize, 0, \
in_feats, kernel, scales, zeros, out_feats, semaphores, num_in_feats, num_out_channels, num_in_channels); in_feats, kernel, scales, zeros, out_feats, semaphores, num_in_feats, num_out_channels, num_in_channels);
template<int N> template<int N>
...@@ -90,8 +91,8 @@ __inline__ __device__ uint32_t cast_smem_ptr_to_uint(void const *const ptr) { ...@@ -90,8 +91,8 @@ __inline__ __device__ uint32_t cast_smem_ptr_to_uint(void const *const ptr) {
template<typename f16_t> template<typename f16_t>
__inline__ __device__ void ldmatrix_m8n8_x4_b16(f16_t *shared_warp, int ax0_0, uint32_t addr) { __inline__ __device__ void ldmatrix_m8n8_x4_b16(f16_t *shared_warp, int ax0_0, uint32_t addr) {
static_assert(std::is_same<f16_t, half>::value || std::is_same<f16_t, __nv_bfloat16>::value, static_assert(std::is_same<f16_t, half>::value || std::is_same<f16_t, __hip_bfloat16>::value,
"ldmatrix_m8n8_x4_b16 supports only half or __nv_bfloat16 types."); "ldmatrix_m8n8_x4_b16 supports only half or __hip_bfloat16 types.");
asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16" asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16"
"{%0, %1, %2, %3}, [%4];" "{%0, %1, %2, %3}, [%4];"
: "=r"(((unsigned *)(shared_warp + (ax0_0 * 8)))[0]), : "=r"(((unsigned *)(shared_warp + (ax0_0 * 8)))[0]),
...@@ -103,8 +104,8 @@ __inline__ __device__ void ldmatrix_m8n8_x4_b16(f16_t *shared_warp, int ax0_0, u ...@@ -103,8 +104,8 @@ __inline__ __device__ void ldmatrix_m8n8_x4_b16(f16_t *shared_warp, int ax0_0, u
template<typename f16_t> template<typename f16_t>
__inline__ __device__ void ldmatrix_m8n8_x4_trans_b16(f16_t *shared_warp, int ax0_0, uint32_t addr) { __inline__ __device__ void ldmatrix_m8n8_x4_trans_b16(f16_t *shared_warp, int ax0_0, uint32_t addr) {
static_assert(std::is_same<f16_t, half>::value || std::is_same<f16_t, __nv_bfloat16>::value, static_assert(std::is_same<f16_t, half>::value || std::is_same<f16_t, __hip_bfloat16>::value,
"ldmatrix_m8n8_x4_trans_b16 supports only half or __nv_bfloat16 types."); "ldmatrix_m8n8_x4_trans_b16 supports only half or __hip_bfloat16 types.");
asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16" asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16"
"{%0, %1, %2, %3}, [%4];" "{%0, %1, %2, %3}, [%4];"
: "=r"(((unsigned *)(shared_warp + (ax0_0 * 8)))[0]), : "=r"(((unsigned *)(shared_warp + (ax0_0 * 8)))[0]),
...@@ -149,7 +150,7 @@ __device__ __inline__ void mma_m16n8k16<half>(float *C_warp, half *A_shared_warp ...@@ -149,7 +150,7 @@ __device__ __inline__ void mma_m16n8k16<half>(float *C_warp, half *A_shared_warp
template<> template<>
__device__ __inline__ void __device__ __inline__ void
mma_m16n8k16<__nv_bfloat16>(float *C_warp, __nv_bfloat16 *A_shared_warp, __nv_bfloat16 *B_shared_warp) { mma_m16n8k16<__hip_bfloat16>(float *C_warp, __hip_bfloat16 *A_shared_warp, __hip_bfloat16 *B_shared_warp) {
asm volatile( asm volatile(
"mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32" "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32"
"{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%10, %11, %12, %13};" "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%10, %11, %12, %13};"
...@@ -378,7 +379,7 @@ __global__ void gemm_w4a16_T1(f16_t *__restrict__ A, ...@@ -378,7 +379,7 @@ __global__ void gemm_w4a16_T1(f16_t *__restrict__ A,
int M, int M,
int N, int N,
int K) { int K) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 #if defined(__DTK_ARCH__) && __DTK_ARCH__ < 800
trap_unsupported_arch(); trap_unsupported_arch();
return; return;
#endif #endif
...@@ -944,7 +945,7 @@ __global__ void gemm_w4a16_T2(f16_t *__restrict__ A, ...@@ -944,7 +945,7 @@ __global__ void gemm_w4a16_T2(f16_t *__restrict__ A,
int M, int M,
int N, int N,
int K) { int K) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 #if defined(__DTK_ARCH__) && __DTK_ARCH__ < 800
trap_unsupported_arch(); trap_unsupported_arch();
return; return;
#endif #endif
...@@ -1277,12 +1278,12 @@ Tensor awq_gemm_forward_cuda(Tensor _in_feats, Tensor _kernel, Tensor _scales, T ...@@ -1277,12 +1278,12 @@ Tensor awq_gemm_forward_cuda(Tensor _in_feats, Tensor _kernel, Tensor _scales, T
dim3 num_blocks((num_out_feats + CTA_M - 1) / CTA_M * j_factors1); dim3 num_blocks((num_out_feats + CTA_M - 1) / CTA_M * j_factors1);
dim3 threads_per_block(WARP_SIZE, NUM_WARPS); dim3 threads_per_block(WARP_SIZE, NUM_WARPS);
auto kernel_func = gemm_w4a16_T2<f16_t, CTA_M, CTA_N, CTA_K, WARP_M, WARP_N, WARP_K, STAGES, G>; auto kernel_func = gemm_w4a16_T2<f16_t, CTA_M, CTA_N, CTA_K, WARP_M, WARP_N, WARP_K, STAGES, G>;
cudaFuncSetAttribute(kernel_func, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemByteSize); hipFuncSetAttribute(kernel_func, hipFuncAttributeMaxDynamicSharedMemorySize, kSmemByteSize);
kernel_func<<<num_blocks, threads_per_block, kSmemByteSize>>>( hipLaunchKernelGGL(( kernel_func), dim3(num_blocks), dim3(threads_per_block), kSmemByteSize, 0,
in_feats, kernel, scales, zeros, out_feats, num_in_feats, num_out_channels, num_in_channels); in_feats, kernel, scales, zeros, out_feats, num_in_feats, num_out_channels, num_in_channels);
} }
} else if (_in_feats.scalar_type() == Tensor::BF16) { } else if (_in_feats.scalar_type() == Tensor::BF16) {
using f16_t = __nv_bfloat16; using f16_t = __hip_bfloat16;
auto in_feats = reinterpret_cast<f16_t *>(_in_feats.data_ptr()); auto in_feats = reinterpret_cast<f16_t *>(_in_feats.data_ptr());
auto kernel = reinterpret_cast<f16_t *>(_kernel.data_ptr<int16_t>()); auto kernel = reinterpret_cast<f16_t *>(_kernel.data_ptr<int16_t>());
...@@ -1357,8 +1358,8 @@ Tensor awq_gemm_forward_cuda(Tensor _in_feats, Tensor _kernel, Tensor _scales, T ...@@ -1357,8 +1358,8 @@ Tensor awq_gemm_forward_cuda(Tensor _in_feats, Tensor _kernel, Tensor _scales, T
dim3 num_blocks((num_out_feats + CTA_M - 1) / CTA_M * j_factors1); dim3 num_blocks((num_out_feats + CTA_M - 1) / CTA_M * j_factors1);
dim3 threads_per_block(WARP_SIZE, NUM_WARPS); dim3 threads_per_block(WARP_SIZE, NUM_WARPS);
auto kernel_func = gemm_w4a16_T2<f16_t, CTA_M, CTA_N, CTA_K, WARP_M, WARP_N, WARP_K, STAGES, G>; auto kernel_func = gemm_w4a16_T2<f16_t, CTA_M, CTA_N, CTA_K, WARP_M, WARP_N, WARP_K, STAGES, G>;
cudaFuncSetAttribute(kernel_func, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemByteSize); hipFuncSetAttribute(kernel_func, hipFuncAttributeMaxDynamicSharedMemorySize, kSmemByteSize);
kernel_func<<<num_blocks, threads_per_block, kSmemByteSize>>>( hipLaunchKernelGGL(( kernel_func), dim3(num_blocks), dim3(threads_per_block), kSmemByteSize, 0,
in_feats, kernel, scales, zeros, out_feats, num_in_feats, num_out_channels, num_in_channels); in_feats, kernel, scales, zeros, out_feats, num_in_feats, num_out_channels, num_in_channels);
} }
} else { } else {
......
src/kernels/awq/gemv_awq.cu → src/kernels/awq/gemv_awq.hip
View file @ 1a8114bf
#include "hip/hip_runtime.h"
/* /*
* Modified from NVIDIA * Modified from NVIDIA
* [TRT-LLM](https://github.com/NVIDIA/TensorRT-LLM/tree/d37b507f41a87457fe9f10f7459d08f5db235745/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv) * [TRT-LLM](https://github.com/NVIDIA/TensorRT-LLM/tree/d37b507f41a87457fe9f10f7459d08f5db235745/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv)
...@@ -30,8 +31,8 @@ ...@@ -30,8 +31,8 @@
#include "../utils.cuh" #include "../utils.cuh"
#include <cuda_fp16.h> #include <hip/hip_fp16.h>
#include <cuda_bf16.h> #include <hip/hip_bf16.h>
#include <stdio.h> #include <stdio.h>
#include "dequantize.cuh" #include "dequantize.cuh"
...@@ -80,7 +81,7 @@ __device__ __forceinline__ packed_as<half, 2>::type half2half2<half>(half x) { ...@@ -80,7 +81,7 @@ __device__ __forceinline__ packed_as<half, 2>::type half2half2<half>(half x) {
} }
template<> template<>
__device__ __forceinline__ packed_as<__nv_bfloat16, 2>::type half2half2<__nv_bfloat16>(__nv_bfloat16 x) { __device__ __forceinline__ packed_as<__hip_bfloat16, 2>::type half2half2<__hip_bfloat16>(__hip_bfloat16 x) {
return __bfloat162bfloat162(x); return __bfloat162bfloat162(x);
} }
...@@ -93,7 +94,7 @@ __device__ __forceinline__ float2 half22float2<half2>(half2 val) { ...@@ -93,7 +94,7 @@ __device__ __forceinline__ float2 half22float2<half2>(half2 val) {
} }
template<> template<>
__device__ __forceinline__ float2 half22float2<__nv_bfloat162>(__nv_bfloat162 val) { __device__ __forceinline__ float2 half22float2<__hip_bfloat162>(__hip_bfloat162 val) {
return __bfloat1622float2(val); return __bfloat1622float2(val);
} }
...@@ -106,8 +107,8 @@ __global__ void gemv_kernel(const half_t *inputs, ...@@ -106,8 +107,8 @@ __global__ void gemv_kernel(const half_t *inputs,
const int IC, const int IC,
const int OC) { const int OC) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 #if defined(__DTK_ARCH__) && __DTK_ARCH__ < 800
if constexpr (std::is_same_v<half_t, __nv_bfloat16>) { if constexpr (std::is_same_v<half_t, __hip_bfloat16>) {
trap_unsupported_arch(); trap_unsupported_arch();
return; return;
} }
...@@ -282,10 +283,10 @@ Tensor gemv_awq( ...@@ -282,10 +283,10 @@ Tensor gemv_awq(
return; return;
} }
if constexpr (M > 0) { if constexpr (M > 0) {
gemv_kernel<half_t, N_PER_BLOCK, M, BLOCK_SIZE, GROUP_SIZE> hipLaunchKernelGGL(( gemv_kernel<half_t, N_PER_BLOCK, M, BLOCK_SIZE, GROUP_SIZE>)
<<<num_blocks, num_threads, 0, getCurrentCUDAStream()>>>( , dim3(num_blocks), dim3(num_threads), 0, getCurrentHIPStreamMasqueradingAsCUDA(),
in_feats, kernel, scaling_factors, zeros, out_feats, k, n); in_feats, kernel, scaling_factors, zeros, out_feats, k, n);
checkCUDA(cudaGetLastError()); checkCUDA(hipGetLastError());
} }
}); });
......
src/kernels/awq/semaphore.h
View file @ 1a8114bf
#include "hip/hip_runtime.h"
/*************************************************************************************************** /***************************************************************************************************
* Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause * SPDX-License-Identifier: BSD-3-Clause
...@@ -55,7 +56,7 @@ public: ...@@ -55,7 +56,7 @@ public:
/// Permit fetching the synchronization mechanism early /// Permit fetching the synchronization mechanism early
__device__ void fetch() { __device__ void fetch() {
if (wait_thread) { if (wait_thread) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700 #if defined(__DTK_ARCH__) && __DTK_ARCH__ >= 700
asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" : "=r"(state) : "l"(lock)); asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" : "=r"(state) : "l"(lock));
#else #else
asm volatile("ld.global.cg.b32 %0, [%1];\n" : "=r"(state) : "l"(lock)); asm volatile("ld.global.cg.b32 %0, [%1];\n" : "=r"(state) : "l"(lock));
...@@ -82,7 +83,7 @@ public: ...@@ -82,7 +83,7 @@ public:
__syncthreads(); __syncthreads();
if (wait_thread) { if (wait_thread) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700 #if defined(__DTK_ARCH__) && __DTK_ARCH__ >= 700
asm volatile("st.global.release.gpu.b32 [%0], %1;\n" : : "l"(lock), "r"(status)); asm volatile("st.global.release.gpu.b32 [%0], %1;\n" : : "l"(lock), "r"(status));
#else #else
asm volatile("st.global.cg.b32 [%0], %1;\n" : : "l"(lock), "r"(status)); asm volatile("st.global.cg.b32 [%0], %1;\n" : : "l"(lock), "r"(status));
......
src/kernels/dispatch_utils.h
View file @ 1a8114bf
...@@ -2,13 +2,13 @@ ...@@ -2,13 +2,13 @@
#include "common.h" #include "common.h"
#include "Tensor.h" #include "Tensor.h"
#include <cuda_fp16.h> #include <hip/hip_fp16.h>
template<typename F> template<typename F>
inline auto dispatchFloat(Tensor::ScalarType scalarType, F &&func) { inline auto dispatchFloat(Tensor::ScalarType scalarType, F &&func) {
switch (scalarType) { switch (scalarType) {
case Tensor::BF16: case Tensor::BF16:
return func.template operator()<__nv_bfloat16>(); return func.template operator()<__hip_bfloat16>();
case Tensor::FP16: case Tensor::FP16:
return func.template operator()<half>(); return func.template operator()<half>();
case Tensor::FP32: case Tensor::FP32:
...@@ -23,7 +23,7 @@ template<typename F> ...@@ -23,7 +23,7 @@ template<typename F>
inline auto dispatchFloat16(Tensor::ScalarType scalarType, F &&func) { inline auto dispatchFloat16(Tensor::ScalarType scalarType, F &&func) {
switch (scalarType) { switch (scalarType) {
case Tensor::BF16: case Tensor::BF16:
return func.template operator()<__nv_bfloat16>(); return func.template operator()<__hip_bfloat16>();
case Tensor::FP16: case Tensor::FP16:
return func.template operator()<half>(); return func.template operator()<half>();
default: default:
...@@ -36,7 +36,7 @@ template<typename F> ...@@ -36,7 +36,7 @@ template<typename F>
inline auto dispatch(Tensor::ScalarType scalarType, F &&func) { inline auto dispatch(Tensor::ScalarType scalarType, F &&func) {
switch (scalarType) { switch (scalarType) {
case Tensor::BF16: case Tensor::BF16:
return func.template operator()<__nv_bfloat16>(); return func.template operator()<__hip_bfloat16>();
case Tensor::FP16: case Tensor::FP16:
return func.template operator()<half>(); return func.template operator()<half>();
case Tensor::FP32: case Tensor::FP32:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment