diff --git a/CMakeLists.txt b/CMakeLists.txt index 7d20ebea7..f6cdd2d07 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -604,6 +604,7 @@ add_library(vllm STATIC src/vllm/model_executor/models/mistral_weights.cpp src/vllm/model_executor/models/gemma4_registry.cpp src/vllm/model_executor/models/gemma4_weights.cpp + src/vllm/model_executor/models/gemma4_moe.cpp src/vllm/model_executor/models/gemma4.cpp src/vllm/model_executor/models/gemma4_mm.cpp src/vllm/model_executor/models/gemma3_registry.cpp @@ -1080,11 +1081,21 @@ if(VLLM_CPP_HIP) src/vllm/platforms/rocm.cpp src/vt/rocm/rocm_backend.hip src/vt/rocm/rocm_rmsnorm.hip + src/vt/rocm/rocm_embedding.hip + src/vt/rocm/rocm_dense_basic.hip + src/vt/rocm/rocm_matmul_hipblaslt.hip + src/vt/rocm/rocm_paged_attn.hip + src/vt/rocm/rocm_gemma4_experts.hip src/vt/rocm/rocm_ops.hip) if(VLLM_CPP_HIP_ARCHITECTURES) set_source_files_properties( src/vt/rocm/rocm_backend.hip src/vt/rocm/rocm_rmsnorm.hip + src/vt/rocm/rocm_embedding.hip + src/vt/rocm/rocm_dense_basic.hip + src/vt/rocm/rocm_matmul_hipblaslt.hip + src/vt/rocm/rocm_paged_attn.hip + src/vt/rocm/rocm_gemma4_experts.hip src/vt/rocm/rocm_ops.hip PROPERTIES HIP_ARCHITECTURES "${VLLM_CPP_HIP_ARCHITECTURES}") endif() @@ -1097,6 +1108,27 @@ if(VLLM_CPP_HIP) else() target_link_libraries(vllm PUBLIC amdhip64) endif() + # hipBLAS + hipBLASLt — dense GEMM (ROCm math libraries). + find_path(VLLM_CPP_HIPBLAS_INCLUDE hipblas/hipblas.h + PATHS ${ROCM_PATH}/include /opt/rocm/include) + find_library(VLLM_CPP_HIPBLAS_LIB NAMES hipblas + PATHS ${ROCM_PATH}/lib ${ROCM_PATH}/lib64 /opt/rocm/lib /opt/rocm/lib64) + find_path(VLLM_CPP_HIPBLASLT_INCLUDE hipblaslt/hipblaslt.h + PATHS ${ROCM_PATH}/include /opt/rocm/include) + find_library(VLLM_CPP_HIPBLASLT_LIB NAMES hipblaslt + PATHS ${ROCM_PATH}/lib ${ROCM_PATH}/lib64 /opt/rocm/lib /opt/rocm/lib64) + if(VLLM_CPP_HIPBLAS_INCLUDE AND VLLM_CPP_HIPBLAS_LIB) + target_include_directories(vllm PUBLIC ${VLLM_CPP_HIPBLAS_INCLUDE}) + target_link_libraries(vllm PUBLIC ${VLLM_CPP_HIPBLAS_LIB}) + message(STATUS "ROCm hipBLAS: ${VLLM_CPP_HIPBLAS_LIB}") + else() + message(FATAL_ERROR "VLLM_CPP_HIP=ON requires hipBLAS (libhipblas + hipblas/hipblas.h).") + endif() + if(VLLM_CPP_HIPBLASLT_INCLUDE AND VLLM_CPP_HIPBLASLT_LIB) + target_include_directories(vllm PUBLIC ${VLLM_CPP_HIPBLASLT_INCLUDE}) + target_link_libraries(vllm PUBLIC ${VLLM_CPP_HIPBLASLT_LIB}) + message(STATUS "ROCm hipBLASLt: ${VLLM_CPP_HIPBLASLT_LIB}") + endif() else() # BIT-ROT GUARD (BACKEND-ROCM W0). The ROCm platform leg is deliberately plain # C++ with no HIP header, which means it CAN be compiled on a machine with no diff --git a/include/vllm/model_executor/model_loader/nvfp4_dequant.h b/include/vllm/model_executor/model_loader/nvfp4_dequant.h index a03a89455..680c26e48 100644 --- a/include/vllm/model_executor/model_loader/nvfp4_dequant.h +++ b/include/vllm/model_executor/model_loader/nvfp4_dequant.h @@ -76,4 +76,12 @@ void DequantNvfp4ToBf16(const uint8_t* packed, const uint8_t* weight_scale_fp8, void DequantFp8ToBf16(const uint8_t* weight_f8, float weight_scale, int64_t numel, uint16_t* out_bf16); +// Channel-wise FP8 (compressed-tensors / llm-compressor FP8_DYNAMIC weights): +// weight_f8 [N, K] F8_E4M3 +// scale_bf16 [N] or [N,1] bf16 per-output-channel scale +// out_bf16 [N, K] +// out[n,k] = bf16( f8(w[n,k]) * bf16_to_f32(scale[n]) ) +void DequantFp8ChannelToBf16(const uint8_t* weight_f8, const uint16_t* scale_bf16, + int64_t N, int64_t K, uint16_t* out_bf16); + } // namespace vllm diff --git a/include/vllm/model_executor/models/gemma4.h b/include/vllm/model_executor/models/gemma4.h index 8321cec93..a5e7beee6 100644 --- a/include/vllm/model_executor/models/gemma4.h +++ b/include/vllm/model_executor/models/gemma4.h @@ -53,6 +53,7 @@ #include "vllm/model_executor/models/model_registry.h" #include "vllm/model_executor/models/qwen3_5.h" // PagedKvCache, ForwardLogits #include "vllm/model_executor/models/qwen3_5_weights.h" // OwnedTensor +#include "vllm/model_executor/models/gemma4_moe.h" // MoE AWQ layer weights #include "vllm/transformers_utils/hf_config.h" #include "vllm/v1/attention/backend.h" // CommonAttentionMetadata #include "vllm/v1/kv_cache_interface.h" @@ -97,9 +98,13 @@ struct Gemma4LayerWeights { OwnedTensor layer_scalar; // bf16 [1] (learned per-layer scalar) Gemma4AttnWeights attn; Gemma4MlpWeights mlp; + // Parallel MoE (26B-A4B): empty when enable_moe_block=false (12B dense). + Gemma4MoeLayerWeights moe; bool is_full_attention = false; // layer_type == "full_attention" bool is_kv_shared = false; // layer_idx >= num_layers - num_kv_shared_layers + bool k_eq_v = false; // no v_proj; V shares K (attention_k_eq_v) int64_t head_dim = 0; // 512 full / 256 sliding + int64_t num_kv_heads = 0; // may differ full vs sliding (global vs local GQA) int64_t kv_target_layer = -1; // for shared layers: source of K/V (-1 = self) }; @@ -115,15 +120,19 @@ struct Gemma4Weights { OwnedTensor final_norm; // bf16 [H] (model.norm, plain RMSNorm) OwnedTensor lm_head; // bf16 [H, vocab] Matmul-B; EMPTY when tied std::vector layers; + // Keeps safetensors mmaps alive for borrowed fused expert tensors (26B MoE). + std::shared_ptr shards_keepalive; }; -// Load the `Gemma4ForConditionalGeneration` text backbone (unsloth/gemma-4-E4B-it) -// safetensors into Gemma4Weights. Name map strips the mm wrapper's -// `model.language_model.` prefix; the vision/audio towers + embed_vision/ -// embed_audio projectors are SKIPPED (handled by the future G2/G3 towers). +// Dense / small loads (shards need not outlive weights). Gemma4Weights LoadGemma4ForConditionalGenerationWeights( const std::vector& shards, const HfConfig& config); +// MoE BF16: experts mmap-borrowed; pass FromSafetensorsOwned shared_ptr. +Gemma4Weights LoadGemma4ForConditionalGenerationWeightsOwned( + std::shared_ptr> shards, + const HfConfig& config); + // The Gemma-4 text backbone forward. Per decoder layer (gemma4.py:709-767): // r=h; input_layernorm(h) -> attn(Q/K/V plain-norm, V weight-less norm, dual // proportional/standard RoPE, per-layer head_dim, YOCO KV-share) -> diff --git a/include/vllm/model_executor/models/gemma4_moe.h b/include/vllm/model_executor/models/gemma4_moe.h new file mode 100644 index 000000000..f288a4a56 --- /dev/null +++ b/include/vllm/model_executor/models/gemma4_moe.h @@ -0,0 +1,81 @@ +// Gemma-4 MoE experts: BF16 fused (Google) or FP8 per-expert (Firworks) + resident. +#pragma once + +#include +#include +#include +#include + +#include "vllm/model_executor/models/qwen3_5_weights.h" +#include "vt/device.h" +#include "vt/tensor.h" + +namespace vllm { + +struct Gemma4Weights; + +// One FP8 expert (compressed-tensors channel scales). Host mmap borrows. +struct Gemma4Fp8ExpertMats { + OwnedTensor gate_w; // F8 as I8 [I,H] + OwnedTensor gate_s; // BF16 [I] or [I,1] + OwnedTensor up_w; + OwnedTensor up_s; + OwnedTensor down_w; // F8 [H,I] + OwnedTensor down_s; // BF16 [H] + // Lazy host BF16 cache after first dequant (decode reuse). + mutable std::vector cached_gu; // [2I,H] + mutable std::vector cached_dn; // [H,I] + // Lazy device BF16 copy on compute GPU (avoids H2D every token). + mutable void* dev_gu = nullptr; // [2I,H] bf16 + mutable void* dev_dn = nullptr; // [H,I] bf16 +}; + +struct Gemma4FusedExperts { + // Google BF16 fused stacks (optional). + OwnedTensor gate_up; // bf16 [E, 2I, H] + OwnedTensor down; // bf16 [E, H, I] + // Firworks FP8 per-expert (optional). size()==E when is_fp8. + bool is_fp8 = false; + std::vector fp8; + int64_t num_experts = 0; + int64_t intermediate = 0; + int64_t hidden = 0; + // Optional device-resident BF16 fused stacks after Prepare. + mutable void* gate_up_dev = nullptr; + mutable void* down_dev = nullptr; + mutable int dev_id = -1; + bool Empty() const { return gate_up.Empty() && fp8.empty(); } +}; + +struct Gemma4MoeLayerWeights { + bool enabled = false; + OwnedTensor router_scale; + OwnedTensor router_proj; + OwnedTensor router_proj_fused; + OwnedTensor per_expert_scale; + OwnedTensor pre_feedforward_layernorm_2; + OwnedTensor post_feedforward_layernorm_1; + OwnedTensor post_feedforward_layernorm_2; + Gemma4FusedExperts experts; + int top_k = 8; + int64_t moe_intermediate = 0; +}; + +struct Gemma4MoeScratch { + vt::Tensor tensor; + std::shared_ptr storage; +}; + +Gemma4MoeScratch RunGemma4Moe(vt::Queue& q, const Gemma4MoeLayerWeights& moe, + const vt::Tensor& router_in, const vt::Tensor& expert_in, + int64_t T, int64_t H, float rms_eps); + +size_t UploadGemma4ExpertsResident(std::vector& layers, + int num_gpus); +size_t UploadGemma4ExpertsResidentForWeights(Gemma4Weights& weights, int num_gpus); + +// Dequant one FP8 expert into host BF16 gate_up[2I,H] and down[H,I] (caller-owned). +void DequantGemma4Fp8ExpertToBf16(const Gemma4Fp8ExpertMats& ex, int64_t I, int64_t H, + uint16_t* gate_up_out, uint16_t* down_out); + +} // namespace vllm diff --git a/src/vllm/model_executor/model_loader/nvfp4_dequant.cpp b/src/vllm/model_executor/model_loader/nvfp4_dequant.cpp index 91d2a0e62..b9780ba94 100644 --- a/src/vllm/model_executor/model_loader/nvfp4_dequant.cpp +++ b/src/vllm/model_executor/model_loader/nvfp4_dequant.cpp @@ -85,4 +85,18 @@ void DequantFp8ToBf16(const uint8_t* weight_f8, float weight_scale, } } +void DequantFp8ChannelToBf16(const uint8_t* weight_f8, const uint16_t* scale_bf16, + int64_t N, int64_t K, uint16_t* out_bf16) { + VT_CHECK(weight_f8 != nullptr && scale_bf16 != nullptr && out_bf16 != nullptr, + "fp8 channel dequant: null"); + VT_CHECK(N > 0 && K > 0, "fp8 channel dequant: dims"); + for (int64_t n = 0; n < N; ++n) { + const float s = vt::BF16ToF32(scale_bf16[n]); + const uint8_t* wr = weight_f8 + n * K; + uint16_t* orow = out_bf16 + n * K; + for (int64_t k = 0; k < K; ++k) + orow[k] = vt::F32ToBF16(F8E4M3ToF32(wr[k]) * s); + } +} + } // namespace vllm diff --git a/src/vllm/model_executor/models/gemma4.cpp b/src/vllm/model_executor/models/gemma4.cpp index 91de460bd..cd9d2c44c 100644 --- a/src/vllm/model_executor/models/gemma4.cpp +++ b/src/vllm/model_executor/models/gemma4.cpp @@ -44,7 +44,8 @@ #include "vllm/model_executor/layers/linear.h" // UnquantizedMlpGateUpGeluMethod seam #include "vllm/model_executor/models/dense_attn_block.h" // Dev/DBuf/glue #include "vllm/model_executor/models/device_pool.h" // Pool -#include "vllm/model_executor/models/qwen3_5_common.h" // HostLogits +#include "vllm/model_executor/models/gemma4_moe.h" +#include "vllm/model_executor/models/qwen3_5_common.h" // HostLogits #include "vt/backend.h" #include "vt/dtype.h" #include "vt/ops.h" @@ -201,7 +202,7 @@ DBuf Gemma4AttnBlock(Dev d, const Gemma4LayerWeights& w, const Gemma4Layout& g, double rope_theta_sliding) { const int64_t H = g.hidden; const int64_t Hq = g.num_q_heads; - const int64_t Hkv = g.num_kv_heads; + const int64_t Hkv = w.num_kv_heads > 0 ? w.num_kv_heads : g.num_kv_heads; const float eps = 1e-6f; // rms_norm_eps (E4B) const int64_t qdim = Hq * Dh, kdim = Hkv * Dh; const DType adt = DType::kBF16; @@ -408,12 +409,14 @@ DBuf ForwardBody(Dev d, const std::vector& token_ids, } // --- Per-Layer Embeddings precompute (gemma4.py:845-898). --- - // ple_emb = embed_tokens_per_layer(ids) * sqrt(ple) -> [T, L*ple] - // ple_proj = (per_layer_model_projection @ tok) * hidden^-0.5 -> [T, L*ple] - // reshape [T,L,ple]; RMSNorm(plain) over ple. - // ple_input = (ple_proj + ple_emb) * rsqrt(2) -> [T, L, ple] - DBuf ple_input(d, DType::kBF16, {T, L, ple}); - { + // Skipped when hidden_size_per_layer_input==0 (google/gemma-4-12B-it dense). + DBuf ple_input(d, DType::kBF16, ple > 0 ? std::vector{T, L, ple} + : std::vector{1, 1, 1}); + if (ple > 0) { + // ple_emb = embed_tokens_per_layer(ids) * sqrt(ple) -> [T, L*ple] + // ple_proj = (per_layer_model_projection @ tok) * hidden^-0.5 -> [T, L*ple] + // reshape [T,L,ple]; RMSNorm(plain) over ple. + // ple_input = (ple_proj + ple_emb) * rsqrt(2) -> [T, L, ple] const int64_t LP = L * ple; DBuf ple_emb(d, DType::kBF16, {T, LP}); { @@ -479,19 +482,42 @@ DBuf ForwardBody(Dev d, const std::vector& token_ids, vt::Add(d.q, h1.t(), attn_n.t(), hidden.t()); // dh2 = pre_feedforward_layernorm(h1); mlp; post_feedforward_layernorm; +h1 + // MoE (26B-A4B): parallel dense MLP + MoE on residual (sglang gemma4_causal). Tensor w_pf = ResidentWeight(d, w.pre_feedforward_layernorm, {H}); DBuf dh2(d, DType::kBF16, {T, H}); vt::RmsNorm(d.q, dh2.t(), h1.t(), w_pf, plain); DBuf mlp = Gemma4MlpBlock(d, w.mlp, H, I, dh2.t(), T); - Tensor w_pff = ResidentWeight(d, w.post_feedforward_layernorm, {H}); - DBuf mlp_n(d, DType::kBF16, {T, H}); - vt::RmsNorm(d.q, mlp_n.t(), mlp.t(), w_pff, plain); + DBuf h2(d, DType::kBF16, {T, H}); - vt::Add(d.q, h2.t(), mlp_n.t(), h1.t()); + if (w.moe.enabled) { + // residual for router = h1 (post-attn residual stream) + DBuf moe_in(d, DType::kBF16, {T, H}); + Tensor w_pf2 = ResidentWeight(d, w.moe.pre_feedforward_layernorm_2, {H}); + vt::RmsNorm(d.q, moe_in.t(), h1.t(), w_pf2, plain); + Gemma4MoeScratch moe_out = + RunGemma4Moe(d.q, w.moe, /*router_in=*/h1.t(), /*expert_in=*/moe_in.t(), T, H, eps); + Tensor w_p1 = ResidentWeight(d, w.moe.post_feedforward_layernorm_1, {H}); + Tensor w_p2 = ResidentWeight(d, w.moe.post_feedforward_layernorm_2, {H}); + DBuf n1(d, DType::kBF16, {T, H}); + DBuf n2(d, DType::kBF16, {T, H}); + vt::RmsNorm(d.q, n1.t(), mlp.t(), w_p1, plain); + vt::RmsNorm(d.q, n2.t(), moe_out.tensor, w_p2, plain); + DBuf sum(d, DType::kBF16, {T, H}); + vt::Add(d.q, sum.t(), n1.t(), n2.t()); + Tensor w_pff = ResidentWeight(d, w.post_feedforward_layernorm, {H}); + DBuf n3(d, DType::kBF16, {T, H}); + vt::RmsNorm(d.q, n3.t(), sum.t(), w_pff, plain); + vt::Add(d.q, h2.t(), n3.t(), h1.t()); + } else { + Tensor w_pff = ResidentWeight(d, w.post_feedforward_layernorm, {H}); + DBuf mlp_n(d, DType::kBF16, {T, H}); + vt::RmsNorm(d.q, mlp_n.t(), mlp.t(), w_pff, plain); + vt::Add(d.q, h2.t(), mlp_n.t(), h1.t()); + } // --- PLE (gemma4.py:753-761): gate = gelu(gate_lin(h2)); gated = gate * // ple_l; contrib = post_per_layer_input_norm(proj(gated)); h2 += contrib. --- - { + if (ple > 0) { Tensor wg = ResidentWeight(d, w.per_layer_input_gate, {ple, H}); DBuf gate_lin(d, DType::kBF16, {T, ple}); vt::MatmulBT(d.q, gate_lin.t(), h2.t(), wg); @@ -523,8 +549,10 @@ DBuf ForwardBody(Dev d, const std::vector& token_ids, } // Per-layer learned scalar (gemma4.py:707,765). - const double scalar = static_cast(ReadBf16Scalar(w.layer_scalar)); - vt::MulScalar(d.q, h2.t(), h2.t(), scalar); + if (!w.layer_scalar.Empty()) { + const double scalar = static_cast(ReadBf16Scalar(w.layer_scalar)); + vt::MulScalar(d.q, h2.t(), h2.t(), scalar); + } hidden = std::move(h2); } diff --git a/src/vllm/model_executor/models/gemma4_moe.cpp b/src/vllm/model_executor/models/gemma4_moe.cpp new file mode 100644 index 000000000..7dcb10ec7 --- /dev/null +++ b/src/vllm/model_executor/models/gemma4_moe.cpp @@ -0,0 +1,267 @@ +// Gemma-4 MoE: BF16 fused or FP8 per-expert + optional device resident. +#include "vllm/model_executor/models/gemma4_moe.h" + +#include +#include +#include +#include +#include + +#include "vllm/model_executor/model_loader/nvfp4_dequant.h" +#include "vllm/model_executor/models/dense_attn_block.h" +#include "vllm/model_executor/models/device_pool.h" +#include "vt/backend.h" +#include "vt/dtype.h" +#include "vt/ops.h" + +namespace vllm { +namespace { + +using dense_attn::DBuf; +using dense_attn::Dev; +using dense_attn::ResidentWeight; +using vt::DType; +using vt::Tensor; + +void ExpertGeGLUHost(Dev d, DBuf& out, const Tensor& x, const uint16_t* gate_up_e, + const uint16_t* down_e, int64_t I, int64_t H) { + const int64_t T = x.shape[0]; + DBuf gate_w(d, DType::kBF16, {I, H}, gate_up_e); + DBuf up_w(d, DType::kBF16, {I, H}, gate_up_e + I * H); + DBuf down_w(d, DType::kBF16, {H, I}, down_e); + DBuf gate(d, DType::kBF16, {T, I}); + DBuf up(d, DType::kBF16, {T, I}); + vt::MatmulBT(d.q, gate.t(), x, gate_w.t()); + vt::MatmulBT(d.q, up.t(), x, up_w.t()); + DBuf gu(d, DType::kBF16, {T, 2 * I}); + const size_t row = static_cast(I) * sizeof(uint16_t); + for (int64_t t = 0; t < T; ++t) { + d.b.Copy(d.q, static_cast(gu.ptr()) + static_cast(t) * 2 * row, + static_cast(gate.ptr()) + static_cast(t) * row, row); + d.b.Copy(d.q, static_cast(gu.ptr()) + static_cast(t) * 2 * row + row, + static_cast(up.ptr()) + static_cast(t) * row, row); + } + DBuf act(d, DType::kBF16, {T, I}); + vt::GeluAndMul(d.q, act.t(), gu.t()); + vt::MatmulBT(d.q, out.t(), act.t(), down_w.t()); +} + +void ExpertGeGLUDevice(Dev d, DBuf& out, const Tensor& x, const uint16_t* gate_up_e, + const uint16_t* down_e, int64_t I, int64_t H) { + const int64_t T = x.shape[0]; + const vt::Device dev = d.q.device; + Tensor gate_w = + Tensor::Contiguous(const_cast(gate_up_e), DType::kBF16, dev, {I, H}); + Tensor up_w = Tensor::Contiguous(const_cast(gate_up_e + I * H), DType::kBF16, + dev, {I, H}); + Tensor down_w = + Tensor::Contiguous(const_cast(down_e), DType::kBF16, dev, {H, I}); + DBuf gate(d, DType::kBF16, {T, I}); + DBuf up(d, DType::kBF16, {T, I}); + vt::MatmulBT(d.q, gate.t(), x, gate_w); + vt::MatmulBT(d.q, up.t(), x, up_w); + DBuf gu(d, DType::kBF16, {T, 2 * I}); + const size_t row = static_cast(I) * sizeof(uint16_t); + for (int64_t t = 0; t < T; ++t) { + d.b.Copy(d.q, static_cast(gu.ptr()) + static_cast(t) * 2 * row, + static_cast(gate.ptr()) + static_cast(t) * row, row); + d.b.Copy(d.q, static_cast(gu.ptr()) + static_cast(t) * 2 * row + row, + static_cast(up.ptr()) + static_cast(t) * row, row); + } + DBuf act(d, DType::kBF16, {T, I}); + vt::GeluAndMul(d.q, act.t(), gu.t()); + vt::MatmulBT(d.q, out.t(), act.t(), down_w); +} + +} // namespace + +// Ensure FP8 expert has BF16 cache filled (idempotent). +void EnsureGemma4Fp8ExpertCached(const Gemma4Fp8ExpertMats& ex, int64_t I, int64_t H) { + if (!ex.cached_gu.empty() && !ex.cached_dn.empty() && + static_cast(ex.cached_gu.size()) == 2 * I * H && + static_cast(ex.cached_dn.size()) == H * I) { + return; + } + ex.cached_gu.resize(static_cast(2 * I * H)); + ex.cached_dn.resize(static_cast(H * I)); + DequantFp8ChannelToBf16(ex.gate_w.bytes.data(), + reinterpret_cast(ex.gate_s.bytes.data()), I, H, + ex.cached_gu.data()); + DequantFp8ChannelToBf16(ex.up_w.bytes.data(), + reinterpret_cast(ex.up_s.bytes.data()), I, H, + ex.cached_gu.data() + I * H); + DequantFp8ChannelToBf16(ex.down_w.bytes.data(), + reinterpret_cast(ex.down_s.bytes.data()), H, I, + ex.cached_dn.data()); +} + +// Host BF16 cache + device upload once (subsequent tokens use device GEMM path). +bool EnsureGemma4Fp8ExpertOnDevice(Dev d, const Gemma4Fp8ExpertMats& ex, int64_t I, + int64_t H) { + EnsureGemma4Fp8ExpertCached(ex, I, H); + if (ex.dev_gu != nullptr && ex.dev_dn != nullptr) return true; + const size_t gu_b = static_cast(2 * I * H) * sizeof(uint16_t); + const size_t dn_b = static_cast(H * I) * sizeof(uint16_t); + void* gu = nullptr; + void* dn = nullptr; + try { + gu = d.b.Alloc(gu_b); + dn = d.b.Alloc(dn_b); + d.b.Copy(d.q, gu, ex.cached_gu.data(), gu_b); + d.b.Copy(d.q, dn, ex.cached_dn.data(), dn_b); + d.b.Synchronize(d.q); + ex.dev_gu = gu; + ex.dev_dn = dn; + return true; + } catch (...) { + if (gu) d.b.Free(gu); + if (dn) d.b.Free(dn); + return false; // fall back to host H2D path + } +} + +void DequantGemma4Fp8ExpertToBf16(const Gemma4Fp8ExpertMats& ex, int64_t I, int64_t H, + uint16_t* gate_up_out, uint16_t* down_out) { + VT_CHECK(gate_up_out && down_out, "fp8 expert dequant null out"); + EnsureGemma4Fp8ExpertCached(ex, I, H); + std::memcpy(gate_up_out, ex.cached_gu.data(), ex.cached_gu.size() * sizeof(uint16_t)); + std::memcpy(down_out, ex.cached_dn.data(), ex.cached_dn.size() * sizeof(uint16_t)); +} + +Gemma4MoeScratch RunGemma4Moe(vt::Queue& q, const Gemma4MoeLayerWeights& moe, + const vt::Tensor& router_in, const vt::Tensor& expert_in, + int64_t T, int64_t H, float rms_eps) { + using dense_attn::DBuf; + using dense_attn::Dev; + using dense_attn::ResidentWeight; + using vt::DType; + using vt::Tensor; + + VT_CHECK(moe.enabled && !moe.experts.Empty(), "gemma4 moe: disabled"); + VT_CHECK(router_in.shape[0] == T && router_in.shape[1] == H, "gemma4 moe: router_in"); + VT_CHECK(expert_in.shape[0] == T && expert_in.shape[1] == H, "gemma4 moe: expert_in"); + const int64_t E = moe.experts.num_experts; + const int64_t I = moe.experts.intermediate; + const int top_k = moe.top_k; + VT_CHECK(E > 0 && I > 0 && top_k > 0 && top_k <= E, "gemma4 moe: dims"); + VT_CHECK(moe.experts.hidden == H, "gemma4 moe: H mismatch"); + + Dev d{vt::GetBackend(q.device.type), q}; + const vt::RmsNormArgs plain{rms_eps, false}; + const int compute_dev = q.device.index; + + DBuf rn(d, DType::kBF16, {T, H}); + { + std::vector ones(static_cast(H), vt::F32ToBF16(1.f)); + DBuf w1(d, DType::kBF16, {H}, ones.data()); + vt::RmsNorm(d.q, rn.t(), router_in, w1.t(), plain); + } + + const OwnedTensor& rproj = + !moe.router_proj_fused.Empty() ? moe.router_proj_fused : moe.router_proj; + VT_CHECK(rproj.HasHostBytes() && rproj.nk && rproj.shape[0] == E && rproj.shape[1] == H, + "gemma4 moe: router proj"); + Tensor wp = ResidentWeight(d, rproj); + DBuf logits(d, DType::kF32, {T, E}); + vt::MatmulBT(d.q, logits.t(), rn.t(), wp); + d.b.Synchronize(d.q); + std::vector hlog(static_cast(T * E)); + d.b.Copy(d.q, hlog.data(), logits.ptr(), hlog.size() * sizeof(float)); + d.b.Synchronize(d.q); + + std::vector hscale(static_cast(E), 1.f); + if (moe.per_expert_scale.HasHostBytes()) { + const auto* pe = reinterpret_cast(moe.per_expert_scale.bytes.data()); + for (int64_t e = 0; e < E; ++e) hscale[static_cast(e)] = vt::BF16ToF32(pe[e]); + } + + const auto& ex = moe.experts; + const int64_t gu_stride = 2 * I * H; + const int64_t dn_stride = H * I; + const bool same_dev = + ex.gate_up_dev != nullptr && ex.down_dev != nullptr && ex.dev_id == compute_dev; + + const auto* gu_host = ex.gate_up.Empty() + ? nullptr + : reinterpret_cast(ex.gate_up.bytes.data()); + const auto* dn_host = + ex.down.Empty() ? nullptr : reinterpret_cast(ex.down.bytes.data()); + + DBuf acc(d, DType::kBF16, {T, H}); + acc.Zero(d); + + for (int64_t t = 0; t < T; ++t) { + std::vector idx(static_cast(E)); + for (int e = 0; e < static_cast(E); ++e) idx[static_cast(e)] = e; + std::partial_sort(idx.begin(), idx.begin() + top_k, idx.end(), [&](int a, int b) { + return hlog[static_cast(t * E + a)] > hlog[static_cast(t * E + b)]; + }); + float mx = hlog[static_cast(t * E + idx[0])]; + std::vector wts(static_cast(top_k)); + float sum = 0.f; + for (int i = 0; i < top_k; ++i) { + wts[static_cast(i)] = + std::exp(hlog[static_cast(t * E + idx[static_cast(i)])] - mx); + sum += wts[static_cast(i)]; + } + for (int i = 0; i < top_k; ++i) + wts[static_cast(i)] = + (wts[static_cast(i)] / sum) * + hscale[static_cast(idx[static_cast(i)])]; + + DBuf xin(d, DType::kBF16, {1, H}); + d.b.Copy(d.q, xin.ptr(), + static_cast(expert_in.data) + + static_cast(t) * static_cast(H) * 2, + static_cast(H) * 2); + + DBuf ysum(d, DType::kBF16, {1, H}); + ysum.Zero(d); + for (int i = 0; i < top_k; ++i) { + const int e = idx[static_cast(i)]; + DBuf y(d, DType::kBF16, {1, H}); + if (same_dev) { + auto* gu = static_cast(ex.gate_up_dev) + + static_cast(e) * gu_stride; + auto* dn = + static_cast(ex.down_dev) + static_cast(e) * dn_stride; + ExpertGeGLUDevice(d, y, xin.t(), gu, dn, I, H); + } else if (ex.is_fp8) { + const auto& fex = ex.fp8[static_cast(e)]; + if (EnsureGemma4Fp8ExpertOnDevice(d, fex, I, H)) { + ExpertGeGLUDevice(d, y, xin.t(), static_cast(fex.dev_gu), + static_cast(fex.dev_dn), I, H); + } else { + EnsureGemma4Fp8ExpertCached(fex, I, H); + ExpertGeGLUHost(d, y, xin.t(), fex.cached_gu.data(), fex.cached_dn.data(), I, H); + } + } else { + ExpertGeGLUHost(d, y, xin.t(), gu_host + static_cast(e) * gu_stride, + dn_host + static_cast(e) * dn_stride, I, H); + } + d.b.Synchronize(d.q); + std::vector hy(static_cast(H)), hs(static_cast(H)); + d.b.Copy(d.q, hy.data(), y.ptr(), hy.size() * 2); + d.b.Copy(d.q, hs.data(), ysum.ptr(), hs.size() * 2); + d.b.Synchronize(d.q); + const float ww = wts[static_cast(i)]; + for (int64_t j = 0; j < H; ++j) + hs[static_cast(j)] = vt::F32ToBF16( + vt::BF16ToF32(hs[static_cast(j)]) + + ww * vt::BF16ToF32(hy[static_cast(j)])); + d.b.Copy(d.q, ysum.ptr(), hs.data(), hs.size() * 2); + } + d.b.Copy(d.q, + static_cast(acc.ptr()) + static_cast(t) * static_cast(H) * 2, + ysum.ptr(), static_cast(H) * 2); + } + + Gemma4MoeScratch r; + r.tensor = acc.t(); + const size_t alloc = acc.alloc_bytes(); + void* p = acc.Release(); + r.storage = std::shared_ptr(p, [alloc](void* q) { Pool().Put(alloc, q); }); + return r; +} + +} // namespace vllm diff --git a/src/vllm/model_executor/models/gemma4_registry.cpp b/src/vllm/model_executor/models/gemma4_registry.cpp index aec6201f2..dc2088922 100644 --- a/src/vllm/model_executor/models/gemma4_registry.cpp +++ b/src/vllm/model_executor/models/gemma4_registry.cpp @@ -25,12 +25,17 @@ #include #include "vllm/model_executor/models/gemma4.h" +#include "vllm/model_executor/models/gemma4_moe.h" #include "vllm/model_executor/models/qwen3_5.h" // ForwardLogits (shared carrier) #include "vllm/model_executor/models/qwen3_5_common.h" // HostLogits #include "vllm/v1/kv_cache_dtype.h" #include "vllm/v1/kv_cache_interface.h" #include "vt/dtype.h" +#include +#include +#include + namespace vllm { namespace { @@ -81,15 +86,25 @@ std::unique_ptr LoadGemma4ForConditionalGeneration( } return std::make_unique( registration, - LoadGemma4ForConditionalGenerationWeights(*source.safetensors, config)); + source.safetensors_owned + ? LoadGemma4ForConditionalGenerationWeightsOwned(source.safetensors_owned, + config) + : LoadGemma4ForConditionalGenerationWeights(*source.safetensors, config)); } void PrepareGemma4ForConditionalGeneration(LoadedModel& model, const HfConfig& config, vt::Queue& queue) { - (void)model; (void)config; (void)queue; + const char* env = std::getenv("VT_GEMMA4_RESIDENT_EXPERTS"); + if (env == nullptr || env[0] != '1') return; + auto& gemma = static_cast(model); + Gemma4Weights& w = const_cast(gemma.weights()); + int ngpu = 2; // Upload clamps to hipGetDeviceCount + if (const char* g = std::getenv("VT_GEMMA4_RESIDENT_GPUS")) + ngpu = std::max(1, std::atoi(g)); + UploadGemma4ExpertsResidentForWeights(w, ngpu); } ForwardLogits ForwardGemma4ForConditionalGeneration( @@ -159,6 +174,13 @@ v1::KVCacheConfig MakeGemma4ForConditionalGenerationKVCache(const HfConfig& conf ? config.raw.at("text_config") : config.raw; const int num_kv_heads = static_cast(config.num_key_value_heads); + const int num_kv_heads_full = [&]() { + if (const auto it = raw.find("num_global_key_value_heads"); + it != raw.end() && it->is_number_integer()) { + return it->get(); + } + return num_kv_heads; + }(); const int head_dim_sliding = static_cast(config.head_dim); // 256 int head_dim_full = head_dim_sliding; // 512 if (const auto it = raw.find("global_head_dim"); @@ -168,11 +190,6 @@ v1::KVCacheConfig MakeGemma4ForConditionalGenerationKVCache(const HfConfig& conf const vt::DType kv_dtype = v1::ResolveKvCacheDType(); const int64_t L = config.num_hidden_layers; - // Per-layer full_attention flag from config.layer_types (gemma4.py:441-489); - // a missing/short array defaults a layer to sliding (head_dim_sliding). This - // mirrors gemma4_weights.cpp::MakeTopo (allocation only needs head_dim; YOCO - // kv_target is handled in the forward, not here — G1c would DEDUP the shared - // layers' unused buffers). std::vector is_full(static_cast(L), false); if (const auto it = raw.find("layer_types"); it != raw.end() && it->is_array()) { @@ -188,14 +205,15 @@ v1::KVCacheConfig MakeGemma4ForConditionalGenerationKVCache(const HfConfig& conf kv.kv_cache_groups.emplace_back( std::vector{"fa"}, std::make_shared( - block_size, num_kv_heads, std::max(head_dim_sliding, head_dim_full), - kv_dtype)); + block_size, std::max(num_kv_heads, num_kv_heads_full), + std::max(head_dim_sliding, head_dim_full), kv_dtype)); kv.per_layer_attn_specs.reserve(static_cast(L)); for (int64_t l = 0; l < L; ++l) { - const int hd = is_full[static_cast(l)] ? head_dim_full - : head_dim_sliding; + const bool full = is_full[static_cast(l)]; + const int hd = full ? head_dim_full : head_dim_sliding; + const int hkv = full ? num_kv_heads_full : num_kv_heads; kv.per_layer_attn_specs.push_back(std::make_shared( - block_size, num_kv_heads, hd, kv_dtype)); + block_size, hkv, hd, kv_dtype)); } return kv; } @@ -214,5 +232,10 @@ std::unique_ptr BorrowGemma4LoadedModel(const Gemma4Weights& weight REGISTER_VLLM_MODEL(gemma4, "Gemma4ForConditionalGeneration", kGemma4Factory, kGemma4Info) +// google/gemma-4-12B-it (and other "unified" HF exports) advertise this arch +// name with model_type gemma4_unified. Same text backbone factory; the weight +// loader tolerates no-PLE dense layouts (hidden_size_per_layer_input==0). +REGISTER_VLLM_MODEL(gemma4_unified, "Gemma4UnifiedForConditionalGeneration", + kGemma4Factory, kGemma4Info) } // namespace vllm diff --git a/src/vllm/model_executor/models/gemma4_weights.cpp b/src/vllm/model_executor/models/gemma4_weights.cpp index 8efa17321..4b406fc18 100644 --- a/src/vllm/model_executor/models/gemma4_weights.cpp +++ b/src/vllm/model_executor/models/gemma4_weights.cpp @@ -1,64 +1,39 @@ -// Weight loader for the `Gemma4ForConditionalGeneration` text backbone -// (unsloth/gemma-4-E4B-it, BF16) — MODEL-GEMMA4 G1. Loads the language_model -// stack into Gemma4Weights (gemma4.h) via the shared dense_weight_loaders.h -// helpers, VERIFIED against the checkpoint's safetensors header (2130 tensors; -// the 336 language_model.* tensors mapped below, mm towers skipped). -// -// Grounding: vllm/model_executor/models/gemma4.py — load_weights (:1625-1728) -// strips the `language_model.` prefix (:1644); packed_modules_mapping -// (:1536-1546) qkv_proj<-[q,k,v]_proj, gate_up_proj<-[gate,up]_proj; -// tie_word_embeddings (:1566-1567) — lm_head aliases embed_tokens, checkpoint has -// no lm_head.weight; skip mm weights audio_tower/vision_tower/embed_audio/ -// embed_vision (:1716-1723). -// -// Name map (unsloth/gemma-4-E4B-it, wrapper prefix `model.language_model.`): -// model.language_model.embed_tokens.weight -> embed_tokens [V,H] -// model.language_model.embed_tokens_per_layer.weight -> embed_tokens_per_layer [V, ple*L] -// model.language_model.per_layer_model_projection.weight-> per_layer_model_projection [ple*L, H] -// model.language_model.per_layer_projection_norm.weight -> per_layer_projection_norm [ple] -// model.language_model.norm.weight -> final_norm [H] -// ...layers.N.input_layernorm.weight -> input_layernorm [H] -// ...layers.N.post_attention_layernorm.weight -> post_attention_ln [H] -// ...layers.N.pre_feedforward_layernorm.weight -> pre_feedforward_ln [H] -// ...layers.N.post_feedforward_layernorm.weight -> post_feedforward_ln [H] -// ...layers.N.post_per_layer_input_norm.weight -> post_per_layer_input_norm [H] -// ...layers.N.per_layer_input_gate.weight -> per_layer_input_gate [ple,H] -// ...layers.N.per_layer_projection.weight -> per_layer_projection [H,ple] -// ...layers.N.layer_scalar -> layer_scalar [1] -// ...layers.N.self_attn.{q,k,v}_proj.weight -> merged qkv_proj (raw-NK) -// ...layers.N.self_attn.o_proj.weight -> o_proj (raw-NK) -// ...layers.N.self_attn.{q,k}_norm.weight -> q_norm/k_norm [Dh] -// ...layers.N.mlp.{gate,up}_proj.weight -> merged gate_up_proj (raw-NK) -// ...layers.N.mlp.down_proj.weight -> down_proj (raw-NK) +// Weight loader: Gemma4 E4B PLE / 12B dense BF16 / 26B-A4B MoE BF16 fused experts. +#include +// Experts are mmap-borrowed (30GB host cannot hold full BF16 MoE). #include "vllm/model_executor/models/gemma4.h" +#include +#include #include #include +#include #include #include #include "vllm/model_executor/model_loader/safetensors_reader.h" +#include "vllm/model_executor/model_loader/nvfp4_dequant.h" #include "vllm/model_executor/models/dense_weight_loaders.h" +#include "vllm/model_executor/models/gemma4_moe.h" #include "vt/dtype.h" +#include +#include + namespace vllm { namespace { using dense_loaders::LoadBf16Direct; using dense_loaders::LoadBf16Transposed; using dense_loaders::LoadMergedBf16RawNK; +using dense_loaders::MakeOwned; -// HfConfig::raw is the FULL config.json (hf_config.cpp:414); the Gemma-4 -// language-model scalars (layer_types, global_head_dim, num_kv_shared_layers, -// tie_word_embeddings) are nested under `text_config` in the mm wrapper. Read -// them through this view — top-level for a plain config, text_config otherwise. const nlohmann::json& TextCfg(const nlohmann::json& raw) { const auto it = raw.find("text_config"); if (it != raw.end() && it->is_object()) return *it; return raw; } - bool RawBool(const nlohmann::json& doc, const char* key, bool fallback) { const auto it = doc.find(key); if (it == doc.end() || it->is_null() || !it->is_boolean()) return fallback; @@ -70,25 +45,20 @@ int64_t RawInt(const nlohmann::json& doc, const char* key, int64_t fallback) { return it->get(); } -// Per-layer attention topology derived from the config (gemma4.py:441-489, -// :572-593). full_attention layers use global_head_dim; the last -// num_kv_shared_layers reuse the last non-shared layer of the same type's K/V. struct LayerTopo { std::vector is_full; std::vector is_shared; std::vector head_dim; - std::vector kv_target; // -1 for self, else source layer index + std::vector kv_target; }; LayerTopo MakeTopo(const HfConfig& cfg) { const nlohmann::json& raw = TextCfg(cfg.raw); const int64_t L = cfg.num_hidden_layers; - const int64_t head_dim_sliding = cfg.head_dim; // text_config head_dim (256) - const int64_t head_dim_full = - RawInt(raw, "global_head_dim", cfg.head_dim); // 512 + const int64_t head_dim_sliding = cfg.head_dim; + const int64_t head_dim_full = RawInt(raw, "global_head_dim", cfg.head_dim); const int64_t num_shared = RawInt(raw, "num_kv_shared_layers", 0); const int64_t first_shared = L - num_shared; - LayerTopo t; t.is_full.assign(static_cast(L), false); const auto it = raw.find("layer_types"); @@ -106,7 +76,6 @@ LayerTopo MakeTopo(const HfConfig& cfg) { t.head_dim[static_cast(l)] = full ? head_dim_full : head_dim_sliding; if (num_shared > 0 && l >= first_shared) { t.is_shared[static_cast(l)] = true; - // Last non-shared layer (< first_shared) of the same attention type. int64_t target = -1; for (int64_t p = first_shared - 1; p >= 0; --p) { if (t.is_full[static_cast(p)] == full) { @@ -120,8 +89,155 @@ LayerTopo MakeTopo(const HfConfig& cfg) { return t; } -Gemma4LayerWeights LoadGemma4Layer(const TensorResolver& get, int64_t layer, - const LayerTopo& topo) { +// Borrow BF16 tensor from mmap (no host copy). Keepalive = shards shared_ptr. +OwnedTensor BorrowBf16(const StTensor& t, std::shared_ptr owner, + const std::vector& shape_override = {}) { + VT_CHECK(t.dtype == "BF16", "gemma4: expected BF16 borrow"); + std::vector shape = shape_override.empty() ? t.shape : shape_override; + OwnedTensor o; + o.dtype = vt::DType::kBF16; + o.rank = static_cast(shape.size()); + VT_CHECK(o.rank <= vt::kMaxRank, "gemma4: rank"); + int64_t n = 1; + for (int i = 0; i < o.rank; ++i) { + o.shape[i] = shape[static_cast(i)]; + n *= shape[static_cast(i)]; + } + VT_CHECK(static_cast(n) * 2 == t.nbytes, "gemma4: borrow size mismatch"); + o.bytes = OwnedBytes::Borrow(t.data, t.nbytes, std::move(owner)); + return o; +} + +OwnedTensor BorrowBytes(const StTensor& t, vt::DType dt, size_t elem_size, + std::shared_ptr owner) { + OwnedTensor o; + o.dtype = dt; + o.rank = static_cast(t.shape.size()); + VT_CHECK(o.rank <= vt::kMaxRank, "gemma4: rank"); + int64_t n = 1; + for (int i = 0; i < o.rank; ++i) { + o.shape[i] = t.shape[static_cast(i)]; + n *= t.shape[static_cast(i)]; + } + VT_CHECK(static_cast(n) * elem_size == t.nbytes, "gemma4: borrow bytes size"); + o.bytes = OwnedBytes::Borrow(t.data, t.nbytes, std::move(owner)); + return o; +} + +// FP8 + channel scale → owned BF16 raw-NK [N,K] +OwnedTensor LoadFp8ChannelToBf16RawNk(const TensorResolver& get, const std::string& base) { + const StTensor& w = get(base + ".weight"); + const StTensor& s = get(base + ".weight_scale"); + VT_CHECK(w.dtype == "F8_E4M3", "gemma4 fp8: expected F8_E4M3 for " + base); + VT_CHECK(s.dtype == "BF16", "gemma4 fp8: expected BF16 scale for " + base); + VT_CHECK(w.shape.size() == 2, "gemma4 fp8: rank-2 weight"); + const int64_t N = w.shape[0], K = w.shape[1]; + VT_CHECK(s.shape[0] == N, "gemma4 fp8: scale N"); + OwnedTensor o = MakeOwned(vt::DType::kBF16, {N, K}); + o.nk = true; + DequantFp8ChannelToBf16(w.data, reinterpret_cast(s.data), N, K, + reinterpret_cast(o.bytes.data())); + MaybeReleaseSourcePages(w.data, w.nbytes); + MaybeReleaseSourcePages(s.data, s.nbytes); + return o; +} + +void FuseRouter(Gemma4MoeLayerWeights& m, int64_t E, int64_t H) { + VT_CHECK(m.router_proj.nk && m.router_proj.shape[0] == E && m.router_proj.shape[1] == H, + "gemma4 moe: router_proj shape"); + const float rsqrt_h = 1.f / std::sqrt(static_cast(H)); + const auto* sc = reinterpret_cast(m.router_scale.bytes.data()); + const auto* ps = reinterpret_cast(m.router_proj.bytes.data()); + m.router_proj_fused = MakeOwned(vt::DType::kBF16, {E, H}); + m.router_proj_fused.nk = true; + auto* pd = reinterpret_cast(m.router_proj_fused.bytes.data()); + for (int64_t e = 0; e < E; ++e) + for (int64_t j = 0; j < H; ++j) + pd[e * H + j] = + vt::F32ToBF16(vt::BF16ToF32(ps[e * H + j]) * vt::BF16ToF32(sc[j]) * rsqrt_h); +} + +Gemma4MoeLayerWeights LoadMoeCommonRouter(const TensorResolver& get, const std::string& base, + int64_t /*E*/, int64_t top_k, int64_t moe_I) { + Gemma4MoeLayerWeights m; + m.enabled = true; + m.top_k = static_cast(top_k); + m.moe_intermediate = moe_I; + m.router_scale = LoadBf16Direct(get, base + "router.scale"); + m.router_proj = LoadMergedBf16RawNK(get, {base + "router.proj.weight"}); + m.per_expert_scale = LoadBf16Direct(get, base + "router.per_expert_scale"); + m.pre_feedforward_layernorm_2 = + LoadBf16Direct(get, base + "pre_feedforward_layernorm_2.weight"); + m.post_feedforward_layernorm_1 = + LoadBf16Direct(get, base + "post_feedforward_layernorm_1.weight"); + m.post_feedforward_layernorm_2 = + LoadBf16Direct(get, base + "post_feedforward_layernorm_2.weight"); + return m; +} + +Gemma4MoeLayerWeights LoadMoeBf16( + const TensorResolver& get, const std::string& base, int64_t E, int64_t top_k, + int64_t moe_I, int64_t H, std::shared_ptr owner) { + Gemma4MoeLayerWeights m = LoadMoeCommonRouter(get, base, E, top_k, moe_I); + FuseRouter(m, E, H); + const StTensor& gu = get(base + "experts.gate_up_proj"); + const StTensor& dn = get(base + "experts.down_proj"); + VT_CHECK(gu.shape.size() == 3 && gu.shape[0] == E && gu.shape[1] == 2 * moe_I && + gu.shape[2] == H, + "gemma4 moe: gate_up expected [E,2I,H]"); + VT_CHECK(dn.shape.size() == 3 && dn.shape[0] == E && dn.shape[1] == H && + dn.shape[2] == moe_I, + "gemma4 moe: down expected [E,H,I]"); + m.experts.gate_up = BorrowBf16(gu, owner); + m.experts.down = BorrowBf16(dn, owner); + m.experts.is_fp8 = false; + m.experts.num_experts = E; + m.experts.intermediate = moe_I; + m.experts.hidden = H; + return m; +} + +// Firworks / per-expert FP8 MoELinear export. +Gemma4MoeLayerWeights LoadMoeFp8PerExpert( + const TensorResolver& get, const std::string& base, int64_t E, int64_t top_k, + int64_t moe_I, int64_t H, std::shared_ptr owner) { + Gemma4MoeLayerWeights m = LoadMoeCommonRouter(get, base, E, top_k, moe_I); + FuseRouter(m, E, H); + m.experts.is_fp8 = true; + m.experts.num_experts = E; + m.experts.intermediate = moe_I; + m.experts.hidden = H; + m.experts.fp8.resize(static_cast(E)); + for (int64_t e = 0; e < E; ++e) { + const std::string eb = base + "experts." + std::to_string(e) + "."; + auto& ex = m.experts.fp8[static_cast(e)]; + const StTensor& gw = get(eb + "gate_proj.weight"); + const StTensor& gs = get(eb + "gate_proj.weight_scale"); + const StTensor& uw = get(eb + "up_proj.weight"); + const StTensor& us = get(eb + "up_proj.weight_scale"); + const StTensor& dw = get(eb + "down_proj.weight"); + const StTensor& ds = get(eb + "down_proj.weight_scale"); + VT_CHECK(gw.dtype == "F8_E4M3" && gw.shape.size() == 2 && gw.shape[0] == moe_I && + gw.shape[1] == H, + "gemma4 fp8 expert gate shape"); + VT_CHECK(dw.dtype == "F8_E4M3" && dw.shape[0] == H && dw.shape[1] == moe_I, + "gemma4 fp8 expert down shape"); + ex.gate_w = BorrowBytes(gw, vt::DType::kI8, 1, owner); + ex.gate_s = BorrowBytes(gs, vt::DType::kBF16, 2, owner); + ex.up_w = BorrowBytes(uw, vt::DType::kI8, 1, owner); + ex.up_s = BorrowBytes(us, vt::DType::kBF16, 2, owner); + ex.down_w = BorrowBytes(dw, vt::DType::kI8, 1, owner); + ex.down_s = BorrowBytes(ds, vt::DType::kBF16, 2, owner); + } + return m; +} + +Gemma4LayerWeights LoadGemma4Layer( + const TensorResolver& get, const std::unordered_set& names, + int64_t layer, const LayerTopo& topo, bool load_ple, int64_t num_q_heads, + int64_t num_kv_heads_sliding, int64_t num_kv_heads_full, bool enable_moe, + int64_t num_experts, int64_t top_k, int64_t moe_I, int64_t H, + std::shared_ptr owner) { const std::string base = "model.language_model.layers." + std::to_string(layer) + "."; const std::string sa = base + "self_attn."; @@ -132,8 +248,9 @@ Gemma4LayerWeights LoadGemma4Layer(const TensorResolver& get, int64_t layer, w.is_kv_shared = topo.is_shared[static_cast(layer)]; w.head_dim = topo.head_dim[static_cast(layer)]; w.kv_target_layer = topo.kv_target[static_cast(layer)]; + w.num_kv_heads = + w.is_full_attention ? num_kv_heads_full : num_kv_heads_sliding; - // Four PLAIN RMSNorm weights (gemma4.py:632-641). w.input_layernorm = LoadBf16Direct(get, base + "input_layernorm.weight"); w.post_attention_layernorm = LoadBf16Direct(get, base + "post_attention_layernorm.weight"); @@ -142,72 +259,161 @@ Gemma4LayerWeights LoadGemma4Layer(const TensorResolver& get, int64_t layer, w.post_feedforward_layernorm = LoadBf16Direct(get, base + "post_feedforward_layernorm.weight"); - // PLE per-layer components (gemma4.py:680-707). - w.per_layer_input_gate = - LoadMergedBf16RawNK(get, {base + "per_layer_input_gate.weight"}); - w.per_layer_projection = - LoadMergedBf16RawNK(get, {base + "per_layer_projection.weight"}); - w.post_per_layer_input_norm = - LoadBf16Direct(get, base + "post_per_layer_input_norm.weight"); - w.layer_scalar = LoadBf16Direct(get, base + "layer_scalar"); - - // QKVParallelLinear merged (q,k,v output-row order), o_proj. No bias - // (attention_bias=false). Shared layers still carry q/k/v_proj in the - // checkpoint (verified); the forward reuses the target layer's K/V and - // discards this layer's K/V, but we load them for a complete, verifiable map. - w.attn.qkv_proj = LoadMergedBf16RawNK( - get, {sa + "q_proj.weight", sa + "k_proj.weight", sa + "v_proj.weight"}); - w.attn.o_proj = LoadMergedBf16RawNK(get, {sa + "o_proj.weight"}); - // Per-head PLAIN Q/K RMSNorm (gemma4.py:434-435). V-norm is weight-less. + if (load_ple) { + w.per_layer_input_gate = + LoadMergedBf16RawNK(get, {base + "per_layer_input_gate.weight"}); + w.per_layer_projection = + LoadMergedBf16RawNK(get, {base + "per_layer_projection.weight"}); + w.post_per_layer_input_norm = + LoadBf16Direct(get, base + "post_per_layer_input_norm.weight"); + } + if (names.count(base + "layer_scalar")) + w.layer_scalar = LoadBf16Direct(get, base + "layer_scalar"); + + const bool fp8_attn = names.count(sa + "q_proj.weight_scale") > 0; + const std::string v_name = sa + "v_proj.weight"; + w.k_eq_v = names.count(v_name) == 0; + if (fp8_attn) { + OwnedTensor q = LoadFp8ChannelToBf16RawNk(get, sa + "q_proj"); + OwnedTensor k = LoadFp8ChannelToBf16RawNk(get, sa + "k_proj"); + OwnedTensor v = w.k_eq_v ? k : LoadFp8ChannelToBf16RawNk(get, sa + "v_proj"); + const int64_t nq = q.shape[0], nk = k.shape[0], nv = v.shape[0]; + w.attn.qkv_proj = MakeOwned(vt::DType::kBF16, {nq + nk + nv, H}); + w.attn.qkv_proj.nk = true; + auto* dst = reinterpret_cast(w.attn.qkv_proj.bytes.data()); + std::memcpy(dst, q.bytes.data(), q.bytes.size()); + std::memcpy(dst + nq * H, k.bytes.data(), k.bytes.size()); + std::memcpy(dst + (nq + nk) * H, v.bytes.data(), v.bytes.size()); + w.attn.o_proj = LoadFp8ChannelToBf16RawNk(get, sa + "o_proj"); + } else if (w.k_eq_v) { + w.attn.qkv_proj = LoadMergedBf16RawNK( + get, {sa + "q_proj.weight", sa + "k_proj.weight", sa + "k_proj.weight"}); + w.attn.o_proj = LoadMergedBf16RawNK(get, {sa + "o_proj.weight"}); + } else { + w.attn.qkv_proj = LoadMergedBf16RawNK( + get, {sa + "q_proj.weight", sa + "k_proj.weight", sa + "v_proj.weight"}); + w.attn.o_proj = LoadMergedBf16RawNK(get, {sa + "o_proj.weight"}); + } w.attn.q_norm = LoadBf16Direct(get, sa + "q_norm.weight"); w.attn.k_norm = LoadBf16Direct(get, sa + "k_norm.weight"); - // GeGLU MLP (gemma4.py:234-247): merged gate_up (gate,up), then down. - w.mlp.gate_up_proj = - LoadMergedBf16RawNK(get, {mlp + "gate_proj.weight", mlp + "up_proj.weight"}); - w.mlp.down_proj = LoadMergedBf16RawNK(get, {mlp + "down_proj.weight"}); + if (w.attn.qkv_proj.rank >= 1 && num_q_heads > 0 && w.num_kv_heads > 0) { + const int64_t rows = w.attn.qkv_proj.shape[0]; + const int64_t denom = num_q_heads + 2 * w.num_kv_heads; + if (denom > 0 && rows % denom == 0) { + const int64_t dh = rows / denom; + if (dh > 0) w.head_dim = dh; + } + } + + if (names.count(mlp + "gate_proj.weight_scale")) { + OwnedTensor g = LoadFp8ChannelToBf16RawNk(get, mlp + "gate_proj"); + OwnedTensor u = LoadFp8ChannelToBf16RawNk(get, mlp + "up_proj"); + const int64_t Ig = g.shape[0]; + w.mlp.gate_up_proj = MakeOwned(vt::DType::kBF16, {Ig + u.shape[0], H}); + w.mlp.gate_up_proj.nk = true; + auto* dst = reinterpret_cast(w.mlp.gate_up_proj.bytes.data()); + std::memcpy(dst, g.bytes.data(), g.bytes.size()); + std::memcpy(dst + Ig * H, u.bytes.data(), u.bytes.size()); + w.mlp.down_proj = LoadFp8ChannelToBf16RawNk(get, mlp + "down_proj"); + } else { + w.mlp.gate_up_proj = + LoadMergedBf16RawNK(get, {mlp + "gate_proj.weight", mlp + "up_proj.weight"}); + w.mlp.down_proj = LoadMergedBf16RawNK(get, {mlp + "down_proj.weight"}); + } + + if (enable_moe) { + if (names.count(base + "experts.gate_up_proj")) { + w.moe = LoadMoeBf16(get, base, num_experts, top_k, moe_I, H, owner); + } else if (names.count(base + "experts.0.gate_proj.weight")) { + w.moe = LoadMoeFp8PerExpert(get, base, num_experts, top_k, moe_I, H, owner); + } + } return w; } -} // namespace - -Gemma4Weights LoadGemma4ForConditionalGenerationWeights( - const std::vector& shards, const HfConfig& config) { +Gemma4Weights LoadImpl(const std::vector& shards, + const HfConfig& config, + std::shared_ptr shards_owner) { std::unordered_map where; - for (const SafetensorsFile& shard : shards) - for (const std::string& name : shard.Names()) where[name] = &shard; - const TensorResolver get = - [&where](const std::string& name) -> const StTensor& { + std::unordered_set names; + for (const SafetensorsFile& shard : shards) { + for (const std::string& name : shard.Names()) { + where[name] = &shard; + names.insert(name); + } + } + const TensorResolver get = [&where](const std::string& name) -> const StTensor& { auto it = where.find(name); VT_CHECK(it != where.end(), "gemma4: tensor not found: " + name); return it->second->Get(name); }; - VT_CHECK(config.num_hidden_layers > 0, - "gemma4: num_hidden_layers must be positive"); + VT_CHECK(config.num_hidden_layers > 0, "gemma4: num_hidden_layers"); + const nlohmann::json& text = TextCfg(config.raw); + const int64_t H = config.hidden_size; + const int64_t ple_cfg = RawInt(text, "hidden_size_per_layer_input", 0); + const bool load_ple = + ple_cfg > 0 && + names.count("model.language_model.embed_tokens_per_layer.weight") > 0; + const bool enable_moe = + RawBool(text, "enable_moe_block", false) || + names.count("model.language_model.layers.0.experts.gate_up_proj") > 0 || + names.count("model.language_model.layers.0.experts.0.gate_proj.weight") > 0 || + names.count("model.language_model.layers.0.router.proj.weight") > 0; + const int64_t num_experts = RawInt(text, "num_experts", 0); + const int64_t top_k = RawInt(text, "top_k_experts", 8); + const int64_t moe_I = RawInt(text, "moe_intermediate_size", 0); - const LayerTopo topo = MakeTopo(config); + if (enable_moe) { + VT_CHECK(shards_owner != nullptr, + "gemma4 MoE BF16 requires safetensors_owned keepalive for expert mmap"); + } + const LayerTopo topo = MakeTopo(config); Gemma4Weights w; - // Gemma-4 ties embeddings by default (text_config.tie_word_embeddings=true). - w.tie_word_embeddings = RawBool(TextCfg(config.raw), "tie_word_embeddings", true); - + w.tie_word_embeddings = RawBool(text, "tie_word_embeddings", true); w.embed_tokens = LoadBf16Direct(get, "model.language_model.embed_tokens.weight"); - w.embed_tokens_per_layer = - LoadBf16Direct(get, "model.language_model.embed_tokens_per_layer.weight"); - w.per_layer_model_projection = LoadMergedBf16RawNK( - get, {"model.language_model.per_layer_model_projection.weight"}); - w.per_layer_projection_norm = - LoadBf16Direct(get, "model.language_model.per_layer_projection_norm.weight"); + if (load_ple) { + w.embed_tokens_per_layer = + LoadBf16Direct(get, "model.language_model.embed_tokens_per_layer.weight"); + w.per_layer_model_projection = LoadMergedBf16RawNK( + get, {"model.language_model.per_layer_model_projection.weight"}); + w.per_layer_projection_norm = + LoadBf16Direct(get, "model.language_model.per_layer_projection_norm.weight"); + } w.final_norm = LoadBf16Direct(get, "model.language_model.norm.weight"); - if (!w.tie_word_embeddings) { + if (!w.tie_word_embeddings && names.count("lm_head.weight")) w.lm_head = LoadBf16Transposed(get, "lm_head.weight"); - } + // Keep shards alive for borrowed experts + w.shards_keepalive = shards_owner; + + const int64_t Hq = config.num_attention_heads; + const int64_t Hkv_slide = config.num_key_value_heads; + const int64_t Hkv_full = RawInt(text, "num_global_key_value_heads", Hkv_slide); w.layers.reserve(static_cast(config.num_hidden_layers)); - for (int64_t l = 0; l < config.num_hidden_layers; ++l) - w.layers.push_back(LoadGemma4Layer(get, l, topo)); + for (int64_t l = 0; l < config.num_hidden_layers; ++l) { + w.layers.push_back(LoadGemma4Layer(get, names, l, topo, load_ple, Hq, Hkv_slide, + Hkv_full, enable_moe, num_experts, top_k, moe_I, + H, shards_owner)); + } return w; } +} // namespace + +Gemma4Weights LoadGemma4ForConditionalGenerationWeights( + const std::vector& shards, const HfConfig& config) { + // Dense 12B path: no owner needed. MoE will throw if experts present without owner. + return LoadImpl(shards, config, nullptr); +} + +Gemma4Weights LoadGemma4ForConditionalGenerationWeightsOwned( + std::shared_ptr> shards, + const HfConfig& config) { + VT_CHECK(shards != nullptr, "gemma4: null shards"); + return LoadImpl(*shards, config, std::shared_ptr(shards, shards.get())); +} + } // namespace vllm diff --git a/src/vllm/v1/worker/gpu/runner.cpp b/src/vllm/v1/worker/gpu/runner.cpp index a62024b74..02c71267a 100644 --- a/src/vllm/v1/worker/gpu/runner.cpp +++ b/src/vllm/v1/worker/gpu/runner.cpp @@ -85,6 +85,20 @@ static bool AsyncRunnerEnvDefault() { return AsyncRunnerFlagIsOn(std::getenv("VT_ASYNC_RUNNER")); } +// Async scheduling without a device-resident sampled-token mirror races the +// embed against host-side token reads and degenerates greedy decode (token-0 +// loops / wrong continuations). The mirror is CUDA-only today +// (async_device_mirror()). On discrete ROCm (and any non-CUDA GPU queue) keep +// the synchronous path unless/until a HIP mirror lands — lab R9700 2026-08-07: +// VT_ASYNC_SCHED=0 restored first-token parity with CPU; async ON produced "!". +static bool QueueSupportsAsyncInputCombine(const vt::Queue& queue) { +#ifdef VLLM_CPP_CUDA + if (queue.device.type == vt::DeviceType::kCUDA) return true; +#endif + (void)queue; + return false; +} + // GDN step-geometry diagnostic (default OFF). When VT_GDN_DIAG_STEP_LOG=1, each // execute_model step prints the request count and the live/free recurrent-state // slot geometry to std::cerr. Read ONCE (never per-step getenv); bounded to the @@ -319,7 +333,8 @@ GPUModelRunner::GPUModelRunner( // spliced into token_ids_cpu by update_req_spec_token_ids + prepare_inputs, // so force the sync host input path here. Byte-identical for non-spec // (spec_config_ is nullopt there, so this is AsyncRunnerEnvDefault()). - async_input_combine_ = AsyncRunnerEnvDefault() && !spec_config_.has_value(); + async_input_combine_ = AsyncRunnerEnvDefault() && !spec_config_.has_value() && + QueueSupportsAsyncInputCombine(queue_); initialize_kv_cache(kv_cache_config); ModelRegistry::Prepare(*model_, config_, queue_); } @@ -352,7 +367,8 @@ GPUModelRunner::GPUModelRunner( // spliced into token_ids_cpu by update_req_spec_token_ids + prepare_inputs, // so force the sync host input path here. Byte-identical for non-spec // (spec_config_ is nullopt there, so this is AsyncRunnerEnvDefault()). - async_input_combine_ = AsyncRunnerEnvDefault() && !spec_config_.has_value(); + async_input_combine_ = AsyncRunnerEnvDefault() && !spec_config_.has_value() && + QueueSupportsAsyncInputCombine(queue_); initialize_kv_cache(kv_cache_config); ModelRegistry::Prepare(*model_, config_, queue_); } @@ -1544,6 +1560,12 @@ ModelRunnerOutput GPUModelRunner::sample_tokens( static_cast(i)]; out.sampled_token_ids.push_back(toks); + // Lab debug: VT_DEBUG_SAMPLED=1 prints every greedy token id. + if (const char* e = std::getenv("VT_DEBUG_SAMPLED"); + e != nullptr && e[0] == '1' && !toks.empty()) { + std::fprintf(stderr, "vt-debug sampled req=%d tok=%d\n", i, toks.front()); + } + // Write-back: append each sampled token to slot i's token row so it becomes // the input at its position next step. num_tokens_no_spec is the next free // column (== num_prompt + generated-so-far). T0 non-spec decode = 1 token. diff --git a/src/vt/rocm/rocm_backend.hip b/src/vt/rocm/rocm_backend.hip index ab0099b21..8c6213b70 100644 --- a/src/vt/rocm/rocm_backend.hip +++ b/src/vt/rocm/rocm_backend.hip @@ -111,6 +111,19 @@ class RocmBackend final : public Backend { return p; } void Free(void* p) override { Check(hipFree(p), "hipFree"); } + // Page-locked host memory so sampler D2H is a real async copy and the host + // can read sampled token ids (default Backend::AllocPinned is Alloc → + // hipMalloc on this backend, which is DEVICE memory — reading it as host is + // UB and breaks multi-token greedy decode on discrete GPUs). + void* AllocPinned(size_t bytes) override { + void* p = nullptr; + Check(hipHostMalloc(&p, bytes == 0 ? 1 : bytes, hipHostMallocDefault), + "hipHostMalloc"); + return p; + } + void FreePinned(void* p) override { + if (p != nullptr) Check(hipHostFree(p), "hipHostFree"); + } void Memset(Queue& q, void* p, int value, size_t bytes) override { Check(hipMemsetAsync(p, value, bytes, AsStream(q)), "hipMemsetAsync"); } diff --git a/src/vt/rocm/rocm_dense_basic.hip b/src/vt/rocm/rocm_dense_basic.hip new file mode 100644 index 000000000..c815d69b4 --- /dev/null +++ b/src/vt/rocm/rocm_dense_basic.hip @@ -0,0 +1,740 @@ +// ROCm dense basic kernels — correctness-grade ports for discrete R9700 bring-up. +#include +#include + +#include +#include +#include +#include + +#include "vt/ops.h" + +namespace vt::rocm { +namespace { + +constexpr int kBlock = 256; + +inline void Check(hipError_t err, const char* what) { + if (err != hipSuccess) { + throw std::runtime_error(std::string("vt rocm: ") + what + ": " + + hipGetErrorString(err)); + } +} +inline hipStream_t AsStream(const Queue& q) { + return static_cast(q.handle); +} +inline unsigned GridFor(int64_t n) { + if (n <= 0) return 1; + const int64_t g = (n + kBlock - 1) / kBlock; + return static_cast(g > 65535 ? 65535 : g); +} + +__device__ inline float Ld(const float* p, int64_t i) { return p[i]; } +__device__ inline float Ld(const __hip_bfloat16* p, int64_t i) { + return __bfloat162float(p[i]); +} +__device__ inline void St(float* p, int64_t i, float v) { p[i] = v; } +__device__ inline void St(__hip_bfloat16* p, int64_t i, float v) { + p[i] = __float2bfloat16(v); +} + +template +__global__ void MulScalarK(Tout* out, const Tin* x, int64_t n, float s) { + for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < n; + idx += gridDim.x * blockDim.x) + St(out, idx, Ld(x, idx) * s); +} + +template +__global__ void SoftCapK(Tout* out, const Tin* x, int64_t n, float cap) { + for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < n; + idx += gridDim.x * blockDim.x) + St(out, idx, cap * tanhf(Ld(x, idx) / cap)); +} + +template +__global__ void AddK(T* out, const T* a, const T* b, int64_t n) { + for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < n; + idx += gridDim.x * blockDim.x) + St(out, idx, Ld(a, idx) + Ld(b, idx)); +} + +template +__global__ void AddBcastK(T* out, const T* a, const T* b, int64_t rows, int64_t cols) { + const int64_t n = rows * cols; + for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < n; + idx += gridDim.x * blockDim.x) + St(out, idx, Ld(a, idx) + Ld(b, idx % cols)); +} + +template +__global__ void GeluMulK(Tout* out, const Tin* x, int64_t n, int64_t d) { + for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < n; + idx += gridDim.x * blockDim.x) { + const int64_t i = idx / d, j = idx - i * d; + const float g = Ld(x, i * 2 * d + j); + const float up = Ld(x, i * 2 * d + d + j); + const float inner = 0.7978845608028654f * (g + 0.044715f * g * g * g); + St(out, idx, 0.5f * g * (1.0f + tanhf(inner)) * up); + } +} + +template +__global__ void SiluMulK(Tout* out, const Tin* x, int64_t n, int64_t d) { + for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < n; + idx += gridDim.x * blockDim.x) { + const int64_t i = idx / d, j = idx - i * d; + const float g = Ld(x, i * 2 * d + j); + const float up = Ld(x, i * 2 * d + d + j); + St(out, idx, (g / (1.0f + expf(-g))) * up); + } +} + +// Matmul / MatmulBT: see rocm_matmul_hipblaslt.hip (hipBLASLt). + +__device__ inline double L3Scale(double freq, double sf, double lo, double hi, double omax) { + if (!(sf > 0.0)) return freq; + constexpr double kTwoPi = 6.283185307179586; + const double low_w = omax / lo, high_w = omax / hi, wl = kTwoPi / freq; + double smooth = 0.0; + if (lo != hi) smooth = (omax / wl - lo) / (hi - lo); + if (wl < high_w) return freq; + if (wl > low_w) return freq / sf; + return (1.0 - smooth) * freq / sf + smooth * freq; +} + +template +__global__ void RopeCacheK(float* cs, const Tid* pos, int64_t t, int rot, int64_t half, + double base, double sf, double lo, double hi, double omax) { + const int64_t n = t * half; + for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < n; + idx += gridDim.x * blockDim.x) { + const int64_t pair = idx % half, tok = idx / half; + const int64_t p = static_cast(pos[tok]); + double freq = pow(base, -2.0 * static_cast(pair) / static_cast(rot)); + freq = L3Scale(freq, sf, lo, hi, omax); + const double ang = static_cast(p) * freq; + cs[tok * rot + pair] = static_cast(cos(ang)); + cs[tok * rot + half + pair] = static_cast(sin(ang)); + } +} + +template +__global__ void ArgmaxK(int64_t* out, const Tin* logits, int64_t rows, int64_t cols) { + const int64_t r = blockIdx.x; + if (r >= rows) return; + __shared__ float bv[kBlock]; + __shared__ int bi[kBlock]; + float best = -INFINITY; + int arg = 0x7fffffff; // unused lanes must lose the tree reduction + for (int64_t j = threadIdx.x; j < cols; j += kBlock) { + const float v = Ld(logits, r * cols + j); + if (v > best || (v == best && (int)j < arg)) { + best = v; + arg = (int)j; + } + } + bv[threadIdx.x] = best; + bi[threadIdx.x] = arg; + __syncthreads(); + for (int s = kBlock / 2; s > 0; s /= 2) { + if ((int)threadIdx.x < s) { + if (bv[threadIdx.x + s] > bv[threadIdx.x] || + (bv[threadIdx.x + s] == bv[threadIdx.x] && bi[threadIdx.x + s] < bi[threadIdx.x])) { + bv[threadIdx.x] = bv[threadIdx.x + s]; + bi[threadIdx.x] = bi[threadIdx.x + s]; + } + } + __syncthreads(); + } + if (threadIdx.x == 0) out[r] = static_cast(bi[0]); +} + +} // namespace + +void MulScalarKernelRocm(Queue& q, Tensor& out, const Tensor& x, double scalar) { + const int64_t n = x.Numel(); + if (n == 0) return; + const float s = static_cast(scalar); + hipStream_t st = AsStream(q); + if (x.dtype == DType::kBF16 && out.dtype == DType::kBF16) + MulScalarK<__hip_bfloat16, __hip_bfloat16> + <<>>(out.Ptr<__hip_bfloat16>(), x.Ptr<__hip_bfloat16>(), n, s); + else if (x.dtype == DType::kF32 && out.dtype == DType::kF32) + MulScalarK + <<>>(out.Ptr(), x.Ptr(), n, s); + else if (x.dtype == DType::kBF16 && out.dtype == DType::kF32) + MulScalarK<__hip_bfloat16, float> + <<>>(out.Ptr(), x.Ptr<__hip_bfloat16>(), n, s); + else + VT_CHECK(false, "rocm mul_scalar dtype"); + Check(hipGetLastError(), "mul_scalar"); +} + +void SoftCapKernelRocm(Queue& q, Tensor& out, const Tensor& x, double cap) { + const int64_t n = x.Numel(); + if (n == 0) return; + const float c = static_cast(cap); + hipStream_t st = AsStream(q); + if (x.dtype == DType::kBF16 && out.dtype == DType::kBF16) + SoftCapK<__hip_bfloat16, __hip_bfloat16> + <<>>(out.Ptr<__hip_bfloat16>(), x.Ptr<__hip_bfloat16>(), n, c); + else if (x.dtype == DType::kF32 && out.dtype == DType::kF32) + SoftCapK + <<>>(out.Ptr(), x.Ptr(), n, c); + else if (x.dtype == DType::kBF16 && out.dtype == DType::kF32) + SoftCapK<__hip_bfloat16, float> + <<>>(out.Ptr(), x.Ptr<__hip_bfloat16>(), n, c); + else + VT_CHECK(false, "rocm soft_cap dtype"); + Check(hipGetLastError(), "soft_cap"); +} + +// mixed-dtype elementwise / broadcast add — f32 compute like CPU LoadF32 +template +__global__ void AddMixK(To* out, const Ta* a, const Tb* b, int64_t n) { + for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < n; + idx += gridDim.x * blockDim.x) + St(out, idx, Ld(a, idx) + Ld(b, idx)); +} +template +__global__ void AddMixBcastK(To* out, const Ta* a, const Tb* b, int64_t rows, + int64_t cols) { + const int64_t n = rows * cols; + for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < n; + idx += gridDim.x * blockDim.x) + St(out, idx, Ld(a, idx) + Ld(b, idx % cols)); +} + +void AddKernelRocm(Queue& q, Tensor& out, const Tensor& a, const Tensor& b) { + hipStream_t st = AsStream(q); + const int64_t n_elem = a.Numel(); + if (n_elem == 0) return; + + const bool bcast = (a.rank == 2 && b.rank == 1 && a.shape[1] == b.shape[0]); + if (!(a.rank == b.rank || bcast)) VT_CHECK(false, "rocm add shapes"); + + const int64_t rows = bcast ? a.shape[0] : 0; + const int64_t cols = bcast ? a.shape[1] : 0; + const int64_t n = bcast ? rows * cols : n_elem; + + auto launch = [&](auto ta, auto tb, auto to) { + using Ta = decltype(ta); + using Tb = decltype(tb); + using To = decltype(to); + if (bcast) + AddMixBcastK<<>>( + out.Ptr(), a.Ptr(), b.Ptr(), rows, cols); + else + AddMixK<<>>(out.Ptr(), a.Ptr(), + b.Ptr(), n); + }; + + // Same-dtype fast path + mixed f32/bf16 (lm_head: f32 logits + bf16 bias). + if (a.dtype == DType::kBF16 && b.dtype == DType::kBF16 && out.dtype == DType::kBF16) + launch(__hip_bfloat16{}, __hip_bfloat16{}, __hip_bfloat16{}); + else if (a.dtype == DType::kF32 && b.dtype == DType::kF32 && out.dtype == DType::kF32) + launch(float{}, float{}, float{}); + else if (a.dtype == DType::kF32 && b.dtype == DType::kBF16 && out.dtype == DType::kF32) + launch(float{}, __hip_bfloat16{}, float{}); + else if (a.dtype == DType::kBF16 && b.dtype == DType::kF32 && out.dtype == DType::kF32) + launch(__hip_bfloat16{}, float{}, float{}); + else if (a.dtype == DType::kBF16 && b.dtype == DType::kBF16 && out.dtype == DType::kF32) + launch(__hip_bfloat16{}, __hip_bfloat16{}, float{}); + else if (a.dtype == DType::kF32 && b.dtype == DType::kF32 && out.dtype == DType::kBF16) + launch(float{}, float{}, __hip_bfloat16{}); + else + VT_CHECK(false, "rocm add dtype combo"); + Check(hipGetLastError(), "add"); +} + +void GeluAndMulKernelRocm(Queue& q, Tensor& out, const Tensor& x) { + const int64_t d = x.shape[1] / 2, n = x.shape[0] * d; + if (n == 0) return; + hipStream_t st = AsStream(q); + if (x.dtype == DType::kBF16) + GeluMulK<__hip_bfloat16, __hip_bfloat16><<>>( + out.Ptr<__hip_bfloat16>(), x.Ptr<__hip_bfloat16>(), n, d); + else if (x.dtype == DType::kF32) + GeluMulK + <<>>(out.Ptr(), x.Ptr(), n, d); + else + VT_CHECK(false, "rocm gelu_and_mul dtype"); + Check(hipGetLastError(), "gelu_and_mul"); +} + +void SiluAndMulKernelRocm(Queue& q, Tensor& out, const Tensor& x) { + const int64_t d = x.shape[1] / 2, n = x.shape[0] * d; + if (n == 0) return; + hipStream_t st = AsStream(q); + if (x.dtype == DType::kBF16) + SiluMulK<__hip_bfloat16, __hip_bfloat16><<>>( + out.Ptr<__hip_bfloat16>(), x.Ptr<__hip_bfloat16>(), n, d); + else if (x.dtype == DType::kF32) + SiluMulK + <<>>(out.Ptr(), x.Ptr(), n, d); + else + VT_CHECK(false, "rocm silu_and_mul dtype"); + Check(hipGetLastError(), "silu_and_mul"); +} + +// MatmulKernelRocm / MatmulBTKernelRocm live in rocm_matmul_hipblaslt.hip + +void RopeCosSinCacheKernelRocm(Queue& q, Tensor& cos_sin, const Tensor& pos, + const RopeArgs& args) { + const int64_t t = cos_sin.shape[0], half = args.rotary_dim / 2, n = t * half; + if (n == 0) return; + hipStream_t st = AsStream(q); + const double base = args.base, sf = args.llama3_scaling_factor, lo = args.llama3_low_freq_factor, + hi = args.llama3_high_freq_factor, omax = args.llama3_orig_max_position; + if (pos.dtype == DType::kI32) + RopeCacheK<<>>( + cos_sin.Ptr(), pos.Ptr(), t, args.rotary_dim, half, base, sf, lo, hi, + omax); + else + RopeCacheK<<>>( + cos_sin.Ptr(), pos.Ptr(), t, args.rotary_dim, half, base, sf, lo, hi, + omax); + Check(hipGetLastError(), "rope_cos_sin_cache"); +} + +void GreedyArgmaxKernelRocm(Queue& q, Tensor& out, const Tensor& logits) { + const int64_t rows = logits.shape[0], cols = logits.shape[1]; + if (rows == 0) return; + VT_CHECK(out.dtype == DType::kI64 || out.dtype == DType::kI32, + "rocm greedy_argmax: out must be i64 or i32"); + hipStream_t st = AsStream(q); + // Engine + cross-device tests use i64 token ids. + if (out.dtype == DType::kI64) { + if (logits.dtype == DType::kF32) + ArgmaxK<<(rows), kBlock, 0, st>>>( + out.Ptr(), logits.Ptr(), rows, cols); + else if (logits.dtype == DType::kBF16) + ArgmaxK<__hip_bfloat16><<(rows), kBlock, 0, st>>>( + out.Ptr(), logits.Ptr<__hip_bfloat16>(), rows, cols); + else + VT_CHECK(false, "rocm greedy_argmax logits dtype"); + } else { + // rare i32 path — write via i64 kernel into temp is overkill; cast pointer + VT_CHECK(false, "rocm greedy_argmax: i32 out not used; use i64"); + } + Check(hipGetLastError(), "greedy_argmax"); +} + +// ---------------- CastBf16 / CastF32 (support strided packed rows) ---------- +__global__ void CastBf16K(__hip_bfloat16* out, const float* in, int64_t n, int64_t row_size, + int64_t row_stride) { + for (int64_t i = blockIdx.x * blockDim.x + threadIdx.x; i < n; + i += gridDim.x * blockDim.x) { + const int64_t row = i / row_size, col = i - row * row_size; + out[i] = __float2bfloat16(in[row * row_stride + col]); + } +} +__global__ void CastF32K(float* out, const __hip_bfloat16* in, int64_t n, int64_t row_size, + int64_t row_stride) { + for (int64_t i = blockIdx.x * blockDim.x + threadIdx.x; i < n; + i += gridDim.x * blockDim.x) { + const int64_t row = i / row_size, col = i - row * row_size; + out[i] = __bfloat162float(in[row * row_stride + col]); + } +} + +void CastBf16KernelRocm(Queue& q, Tensor& out, const Tensor& in) { + const int64_t n = out.Numel(); + if (n == 0) return; + const int64_t rows = in.shape[0]; + const int64_t row_size = n / rows; + CastBf16K<<>>( + out.Ptr<__hip_bfloat16>(), in.Ptr(), n, row_size, in.stride[0]); + Check(hipGetLastError(), "cast_bf16"); +} + +void CastF32KernelRocm(Queue& q, Tensor& out, const Tensor& in) { + const int64_t n = out.Numel(); + if (n == 0) return; + const int64_t rows = in.shape[0]; + const int64_t row_size = n / rows; + CastF32K<<>>(out.Ptr(), in.Ptr<__hip_bfloat16>(), + n, row_size, in.stride[0]); + Check(hipGetLastError(), "cast_f32"); +} + +// ---------------- QkvSplit ---------------- +template +__global__ void QkvSplitK(T* q_out, T* k_out, T* v_out, const T* qkv, int64_t t, int64_t q_dim, + int64_t k_dim, int64_t v_dim) { + const int64_t total = q_dim + k_dim + v_dim; + const int64_t nq = t * q_dim, nk = t * k_dim, nv = t * v_dim; + int64_t n = nq; + if (nk > n) n = nk; + if (nv > n) n = nv; + for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < n; + idx += gridDim.x * blockDim.x) { + if (idx < nq) { + const int64_t i = idx / q_dim, j = idx % q_dim; + q_out[idx] = qkv[i * total + j]; + } + if (idx < nk) { + const int64_t i = idx / k_dim, j = idx % k_dim; + k_out[idx] = qkv[i * total + q_dim + j]; + } + if (idx < nv) { + const int64_t i = idx / v_dim, j = idx % v_dim; + v_out[idx] = qkv[i * total + q_dim + k_dim + j]; + } + } +} + +void QkvSplitKernelRocm(Queue& q, Tensor& q_out, Tensor& k_out, Tensor& v_out, + const Tensor& qkv) { + const int64_t t = qkv.shape[0]; + if (t == 0) return; + const int64_t q_dim = q_out.Numel() / t, k_dim = k_out.Numel() / t, v_dim = v_out.Numel() / t; + int64_t n = t * q_dim; + if (t * k_dim > n) n = t * k_dim; + if (t * v_dim > n) n = t * v_dim; + hipStream_t st = AsStream(q); + if (qkv.dtype == DType::kBF16) + QkvSplitK<__hip_bfloat16><<>>( + q_out.Ptr<__hip_bfloat16>(), k_out.Ptr<__hip_bfloat16>(), v_out.Ptr<__hip_bfloat16>(), + qkv.Ptr<__hip_bfloat16>(), t, q_dim, k_dim, v_dim); + else if (qkv.dtype == DType::kF32) + QkvSplitK<<>>( + q_out.Ptr(), k_out.Ptr(), v_out.Ptr(), qkv.Ptr(), t, q_dim, + k_dim, v_dim); + else + VT_CHECK(false, "rocm qkv_split dtype"); + Check(hipGetLastError(), "qkv_split"); +} + +// ---------------- LayerNorm (tree reduce — wavefront-agnostic) ------------- +template +__global__ void LayerNormRowK(OutT* out, const InT* x, const WT* weight, const WT* bias, + int64_t d, float eps) { + const int64_t row = blockIdx.x; + const InT* xr = x + row * d; + OutT* orow = out + row * d; + __shared__ float partial[kBlock]; + + float sum = 0.f; + for (int64_t i = threadIdx.x; i < d; i += kBlock) sum += Ld(xr, i); + partial[threadIdx.x] = sum; + __syncthreads(); + for (int s = kBlock / 2; s > 0; s /= 2) { + if ((int)threadIdx.x < s) partial[threadIdx.x] += partial[threadIdx.x + s]; + __syncthreads(); + } + const float mean = partial[0] / static_cast(d); + __syncthreads(); + + float sq = 0.f; + for (int64_t i = threadIdx.x; i < d; i += kBlock) { + const float dv = Ld(xr, i) - mean; + sq += dv * dv; + } + partial[threadIdx.x] = sq; + __syncthreads(); + for (int s = kBlock / 2; s > 0; s /= 2) { + if ((int)threadIdx.x < s) partial[threadIdx.x] += partial[threadIdx.x + s]; + __syncthreads(); + } + const float rstd = rsqrtf(partial[0] / static_cast(d) + eps); + + for (int64_t i = threadIdx.x; i < d; i += kBlock) { + float v = (Ld(xr, i) - mean) * rstd; + if (weight != nullptr) v *= Ld(weight, i); + if (bias != nullptr) v += Ld(bias, i); + St(orow, i, v); + } +} + +void LayerNormKernelRocm(Queue& q, Tensor& out, const Tensor& x, const Tensor* weight, + const Tensor* bias, const LayerNormArgs& args) { + const int64_t d = x.shape[x.rank - 1]; + const int64_t rows = d == 0 ? 0 : x.Numel() / d; + if (rows == 0 || d == 0) return; + hipStream_t st = AsStream(q); + const unsigned grid = static_cast(rows < 65535 ? rows : 65535); + auto launch = [&](auto out_tag, auto in_tag, auto w_tag) { + using OutT = decltype(out_tag); + using InT = decltype(in_tag); + using WT = decltype(w_tag); + const WT* wp = weight ? static_cast(weight->data) : nullptr; + const WT* bp = bias ? static_cast(bias->data) : nullptr; + LayerNormRowK<<>>( + static_cast(out.data), static_cast(x.data), wp, bp, d, args.eps); + }; + const DType wdt = + weight ? weight->dtype : (bias ? bias->dtype : x.dtype); + if (out.dtype == DType::kBF16 && x.dtype == DType::kBF16 && wdt == DType::kBF16) + launch(__hip_bfloat16{}, __hip_bfloat16{}, __hip_bfloat16{}); + else if (out.dtype == DType::kBF16 && x.dtype == DType::kBF16 && wdt == DType::kF32) + launch(__hip_bfloat16{}, __hip_bfloat16{}, float{}); + else if (out.dtype == DType::kF32 && x.dtype == DType::kF32 && wdt == DType::kF32) + launch(float{}, float{}, float{}); + else if (out.dtype == DType::kBF16 && x.dtype == DType::kF32 && wdt == DType::kF32) + launch(__hip_bfloat16{}, float{}, float{}); + else if (out.dtype == DType::kF32 && x.dtype == DType::kBF16 && wdt == DType::kBF16) + launch(float{}, __hip_bfloat16{}, __hip_bfloat16{}); + else if (out.dtype == DType::kF32 && x.dtype == DType::kBF16 && wdt == DType::kF32) + launch(float{}, __hip_bfloat16{}, float{}); + else + VT_CHECK(false, "rocm layer_norm dtype"); + Check(hipGetLastError(), "layer_norm"); +} + +// ---------------- Relu ---------------- +template +__global__ void ReluK(T* out, const T* x, int64_t n) { + for (int64_t i = blockIdx.x * blockDim.x + threadIdx.x; i < n; + i += gridDim.x * blockDim.x) { + const float v = Ld(x, i); + St(out, i, v > 0.f ? v : 0.f); + } +} + +void ReluKernelRocm(Queue& q, Tensor& out, const Tensor& x) { + const int64_t n = x.Numel(); + if (n == 0) return; + hipStream_t st = AsStream(q); + if (x.dtype == DType::kBF16) + ReluK<__hip_bfloat16><<>>(out.Ptr<__hip_bfloat16>(), + x.Ptr<__hip_bfloat16>(), n); + else if (x.dtype == DType::kF32) + ReluK<<>>(out.Ptr(), x.Ptr(), n); + else + VT_CHECK(false, "rocm relu dtype"); + Check(hipGetLastError(), "relu"); +} + +// ---------------- RopeNeox (in-place on qs/ks) ----------------------------- +template +__global__ void RopeNeoxK(T* qs, T* ks, const Tid* pos, int64_t hq, int64_t hk, int64_t d, + int64_t half, int rot, double base, double l3_sf, double l3_lo, + double l3_hi, double l3_omax, int64_t n) { + const int64_t heads = hq + hk; + for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < n; + idx += gridDim.x * blockDim.x) { + const int64_t pair = idx % half; + const int64_t head = (idx / half) % heads; + const int64_t tok = idx / (half * heads); + T* ptr; + int64_t off; + if (head < hq) { + ptr = qs; + off = (tok * hq + head) * d; + } else { + ptr = ks; + off = (tok * hk + (head - hq)) * d; + } + const int64_t p = static_cast(pos[tok]); + double freq = pow(base, -2.0 * static_cast(pair) / static_cast(rot)); + freq = L3Scale(freq, l3_sf, l3_lo, l3_hi, l3_omax); + const double angle = static_cast(p) * freq; + const float c = static_cast(cos(angle)); + const float sn = static_cast(sin(angle)); + const float x = Ld(ptr, off + pair); + const float y = Ld(ptr, off + pair + half); + St(ptr, off + pair, x * c - y * sn); + St(ptr, off + pair + half, x * sn + y * c); + } +} + +void RopeNeoxKernelRocm(Queue& q, Tensor& qs, Tensor& ks, const Tensor& pos, + const RopeArgs& args) { + const int64_t t = qs.shape[0], hq = qs.shape[1], hk = ks.shape[1], d = qs.shape[2]; + const int64_t half = args.rotary_dim / 2; + const int64_t n = t * (hq + hk) * half; + if (n == 0) return; + hipStream_t st = AsStream(q); + const double base = args.base, sf = args.llama3_scaling_factor, lo = args.llama3_low_freq_factor, + hi = args.llama3_high_freq_factor, omax = args.llama3_orig_max_position; + if (qs.dtype == DType::kBF16) { + if (pos.dtype == DType::kI32) + RopeNeoxK<__hip_bfloat16, int32_t><<>>( + qs.Ptr<__hip_bfloat16>(), ks.Ptr<__hip_bfloat16>(), pos.Ptr(), hq, hk, d, + half, args.rotary_dim, base, sf, lo, hi, omax, n); + else + RopeNeoxK<__hip_bfloat16, int64_t><<>>( + qs.Ptr<__hip_bfloat16>(), ks.Ptr<__hip_bfloat16>(), pos.Ptr(), hq, hk, d, + half, args.rotary_dim, base, sf, lo, hi, omax, n); + } else if (qs.dtype == DType::kF32) { + if (pos.dtype == DType::kI32) + RopeNeoxK<<>>( + qs.Ptr(), ks.Ptr(), pos.Ptr(), hq, hk, d, half, args.rotary_dim, + base, sf, lo, hi, omax, n); + else + RopeNeoxK<<>>( + qs.Ptr(), ks.Ptr(), pos.Ptr(), hq, hk, d, half, args.rotary_dim, + base, sf, lo, hi, omax, n); + } else + VT_CHECK(false, "rocm rope_neox dtype"); + Check(hipGetLastError(), "rope_neox"); +} + +// ---------------- RopeFromCache -------------------------------------------- +template +__global__ void RopeFromCacheK(T* qs, T* ks, const Tid* positions, const T* cache, + int64_t cache_rows, int64_t tokens, int64_t hq, int64_t hk, + int64_t q_tok_stride, int64_t q_head_stride, int64_t k_tok_stride, + int64_t k_head_stride, int rotary_dim, int64_t half, + bool is_neox_style, int64_t n) { + const int64_t heads = hq + hk; + for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < n; + idx += gridDim.x * blockDim.x) { + const int64_t pair = idx % half; + const int64_t head = (idx / half) % heads; + const int64_t token = idx / (half * heads); + const int64_t position = static_cast(positions[token]); + if (position < 0 || position >= cache_rows) continue; + const int64_t cache_offset = position * rotary_dim; + const float c = Ld(cache, cache_offset + pair); + const float sn = Ld(cache, cache_offset + half + pair); + T* states = head < hq ? qs : ks; + const int64_t local_head = head < hq ? head : head - hq; + const int64_t row = head < hq ? token * q_tok_stride + local_head * q_head_stride + : token * k_tok_stride + local_head * k_head_stride; + const int64_t first = is_neox_style ? pair : pair * 2; + const int64_t second = is_neox_style ? pair + half : pair * 2 + 1; + const float x = Ld(states, row + first); + const float y = Ld(states, row + second); + St(states, row + first, x * c - y * sn); + St(states, row + second, x * sn + y * c); + } +} + +void RopeFromCacheKernelRocm(Queue& q, Tensor& qs, Tensor* ks, const Tensor& positions, + const Tensor& cache, const RopeArgs& args) { + const int64_t tokens = qs.shape[0]; + const int64_t hq = qs.shape[1]; + const int64_t hk = ks == nullptr ? 0 : ks->shape[1]; + const int64_t half = args.rotary_dim / 2; + const int64_t n = tokens * (hq + hk) * half; + if (n == 0) return; + // mrope (positions.rank==2) not needed for Gemma/Phi dense smoke yet + VT_CHECK(positions.rank == 1, "rocm rope_from_cache: mrope not implemented yet"); + hipStream_t st = AsStream(q); + auto launch = [&](auto t_tag, auto id_tag) { + using T = decltype(t_tag); + using Tid = decltype(id_tag); + RopeFromCacheK<<>>( + qs.Ptr(), ks == nullptr ? nullptr : ks->Ptr(), positions.Ptr(), + cache.Ptr(), cache.shape[0], tokens, hq, hk, qs.stride[0], qs.stride[1], + ks == nullptr ? 0 : ks->stride[0], ks == nullptr ? 0 : ks->stride[1], args.rotary_dim, + half, args.is_neox_style, n); + }; + if (qs.dtype == DType::kBF16) { + if (positions.dtype == DType::kI32) + launch(__hip_bfloat16{}, int32_t{}); + else + launch(__hip_bfloat16{}, int64_t{}); + } else if (qs.dtype == DType::kF32) { + if (positions.dtype == DType::kI32) + launch(float{}, int32_t{}); + else + launch(float{}, int64_t{}); + } else + VT_CHECK(false, "rocm rope_from_cache dtype"); + Check(hipGetLastError(), "rope_from_cache"); +} + +// ---------------- ReshapeAndCache (NHD paged KV write) --------------------- +template +__global__ void ReshapeAndCacheK(const Word* key, const Word* value, Word* key_cache, + Word* value_cache, const int64_t* slot_mapping, + int64_t block_size, int64_t n_elems, int64_t k_block_stride, + int64_t k_page_stride, int64_t v_block_stride, + int64_t v_page_stride, int64_t k_tok_stride, + int64_t v_tok_stride) { + const int64_t token = blockIdx.x; + const int64_t slot = slot_mapping[token]; + if (slot < 0) return; + const int64_t block = slot / block_size; + const int64_t offset = slot % block_size; + const int64_t kdst = block * k_block_stride + offset * k_page_stride; + const int64_t vdst = block * v_block_stride + offset * v_page_stride; + const int64_t ksrc = token * k_tok_stride; + const int64_t vsrc = token * v_tok_stride; + for (int64_t e = threadIdx.x; e < n_elems; e += blockDim.x) { + key_cache[kdst + e] = key[ksrc + e]; + value_cache[vdst + e] = value[vsrc + e]; + } +} + +void ReshapeAndCacheKernelRocm(Queue& q, const Tensor& k, const Tensor& v, Tensor& k_cache, + Tensor& v_cache, const Tensor& slot_mapping) { + const int64_t num_slots = slot_mapping.shape[0]; + const int64_t block_size = k_cache.shape[1]; + const int64_t n_elems = k_cache.shape[2] * k_cache.shape[3]; + if (num_slots == 0 || n_elems == 0) return; + hipStream_t st = AsStream(q); + const unsigned grid = static_cast(num_slots); + const unsigned block = static_cast(n_elems < 512 ? n_elems : 512); + const int64_t* slots = slot_mapping.Ptr(); + switch (SizeOf(k.dtype)) { + case 4: + ReshapeAndCacheK<<>>( + k.Ptr(), v.Ptr(), k_cache.Ptr(), + v_cache.Ptr(), slots, block_size, n_elems, k_cache.stride[0], + k_cache.stride[1], v_cache.stride[0], v_cache.stride[1], k.stride[0], v.stride[0]); + break; + case 2: + ReshapeAndCacheK<<>>( + k.Ptr(), v.Ptr(), k_cache.Ptr(), + v_cache.Ptr(), slots, block_size, n_elems, k_cache.stride[0], + k_cache.stride[1], v_cache.stride[0], v_cache.stride[1], k.stride[0], v.stride[0]); + break; + default: + VT_CHECK(false, "rocm reshape_and_cache: unsupported elem size"); + } + Check(hipGetLastError(), "reshape_and_cache"); +} + +// ---------------- GeluTanh / GeluErf (elementwise) ------------------------- +template +__global__ void GeluTanhK(T* out, const T* x, int64_t n) { + for (int64_t i = blockIdx.x * blockDim.x + threadIdx.x; i < n; + i += gridDim.x * blockDim.x) { + const float v = Ld(x, i); + const float inner = 0.7978845608028654f * (v + 0.044715f * v * v * v); + St(out, i, 0.5f * v * (1.0f + tanhf(inner))); + } +} +template +__global__ void GeluErfK(T* out, const T* x, int64_t n) { + for (int64_t i = blockIdx.x * blockDim.x + threadIdx.x; i < n; + i += gridDim.x * blockDim.x) { + const float v = Ld(x, i); + St(out, i, 0.5f * v * (1.0f + erff(v * 0.7071067811865476f))); + } +} + +void GeluTanhKernelRocm(Queue& q, Tensor& out, const Tensor& x) { + const int64_t n = x.Numel(); + if (n == 0) return; + hipStream_t st = AsStream(q); + if (x.dtype == DType::kBF16) + GeluTanhK<__hip_bfloat16><<>>(out.Ptr<__hip_bfloat16>(), + x.Ptr<__hip_bfloat16>(), n); + else if (x.dtype == DType::kF32) + GeluTanhK<<>>(out.Ptr(), x.Ptr(), n); + else + VT_CHECK(false, "rocm gelu_tanh dtype"); + Check(hipGetLastError(), "gelu_tanh"); +} + +void GeluErfKernelRocm(Queue& q, Tensor& out, const Tensor& x) { + const int64_t n = x.Numel(); + if (n == 0) return; + hipStream_t st = AsStream(q); + if (x.dtype == DType::kBF16) + GeluErfK<__hip_bfloat16><<>>(out.Ptr<__hip_bfloat16>(), + x.Ptr<__hip_bfloat16>(), n); + else if (x.dtype == DType::kF32) + GeluErfK<<>>(out.Ptr(), x.Ptr(), n); + else + VT_CHECK(false, "rocm gelu_erf dtype"); + Check(hipGetLastError(), "gelu_erf"); +} + +} // namespace vt::rocm diff --git a/src/vt/rocm/rocm_embedding.hip b/src/vt/rocm/rocm_embedding.hip new file mode 100644 index 000000000..3e3a30bf2 --- /dev/null +++ b/src/vt/rocm/rocm_embedding.hip @@ -0,0 +1,151 @@ +// ROCm Embedding — second registered kernel (BACKEND-ROCM, M3 start). +// +// PORTED FROM: src/vt/cuda/cuda_ops.cu EmbeddingKernel (grid-stride gather with +// in-kernel id clamp + atomic error flag). Simplified host path: allocate a +// one-shot error flag, launch, synchronize, throw — same correctness contract +// as the original CUDA M0.6 shape (ops.h). The CUDA ring is a later speed +// concern; discrete R9700 needs a working op first. +// +// Wavefront-agnostic: no warp primitives, only grid-stride over elements. +#include +#include + +#include +#include +#include + +#include "vt/ops.h" + +namespace vt::rocm { +namespace { + +constexpr int kBlock = 256; + +inline void Check(hipError_t err, const char* what) { + if (err != hipSuccess) { + throw std::runtime_error(std::string("vt rocm: ") + what + ": " + + hipGetErrorString(err)); + } +} + +inline unsigned GridFor(int64_t n) { + if (n <= 0) return 1; + const int64_t g = (n + kBlock - 1) / kBlock; + return static_cast(g > 65535 ? 65535 : g); +} + +struct EmbeddingErr { + int status; // 0 ok, 1 bad id + int pad; + long long id; +}; + +template +__device__ inline T Load(const T* p, int64_t i) { + return p[i]; +} +template +__device__ inline void Store(T* p, int64_t i, T v) { + p[i] = v; +} + +// bf16 load/store via float only when casting between dtypes — same-dtype path +// copies the 16-bit word. +template +__global__ void EmbeddingKernel(Tout* out, const Tin* table, const Tid* ids, + int64_t n, int64_t h, int64_t v, + EmbeddingErr* err) { + const int64_t step = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + idx < n; idx += step) { + const int64_t i = idx / h; + const int64_t j = idx - i * h; + int64_t id = static_cast(ids[i]); + if (id < 0 || id >= v) { + if (atomicCAS(&err->status, 0, 1) == 0) err->id = static_cast(id); + id = id < 0 ? 0 : v - 1; + } + if constexpr (std::is_same_v) { + Store(out, idx, Load(table, id * h + j)); + } else if constexpr (std::is_same_v && + std::is_same_v) { + Store(out, idx, __float2bfloat16(Load(table, id * h + j))); + } else if constexpr (std::is_same_v && + std::is_same_v) { + Store(out, idx, __bfloat162float(Load(table, id * h + j))); + } else { + Store(out, idx, static_cast(Load(table, id * h + j))); + } + } +} + +template +void LaunchEmbedding(hipStream_t s, Tensor& out, const Tensor& table, + const Tensor& ids, EmbeddingErr* err) { + const int64_t t = ids.shape[0], h = table.shape[1], v = table.shape[0]; + const int64_t n = t * h; + if (n == 0) return; + if (ids.dtype == DType::kI32) { + EmbeddingKernel<<>>( + out.Ptr(), table.Ptr(), ids.Ptr(), n, h, v, err); + } else { + EmbeddingKernel<<>>( + out.Ptr(), table.Ptr(), ids.Ptr(), n, h, v, err); + } +} + +template +void LaunchEmbeddingIn(hipStream_t s, Tensor& out, const Tensor& table, + const Tensor& ids, EmbeddingErr* err) { + if (out.dtype == DType::kF32) { + LaunchEmbedding(s, out, table, ids, err); + } else { + LaunchEmbedding(s, out, table, ids, err); + } +} + +} // namespace + +void EmbeddingKernelRocm(Queue& q, Tensor& out, const Tensor& table, + const Tensor& ids) { + VT_CHECK(table.dtype == DType::kF32 || table.dtype == DType::kBF16, + "rocm embedding: unsupported table dtype (f32/bf16 only)"); + VT_CHECK(out.dtype == DType::kF32 || out.dtype == DType::kBF16, + "rocm embedding: unsupported out dtype"); + VT_CHECK(ids.dtype == DType::kI32 || ids.dtype == DType::kI64, + "rocm embedding: ids must be i32 or i64"); + VT_CHECK(table.rank == 2 && out.rank == 2 && ids.rank == 1, + "rocm embedding: expected table[V,H], out[T,H], ids[T]"); + VT_CHECK(out.shape[0] == ids.shape[0] && out.shape[1] == table.shape[1], + "rocm embedding: shape mismatch"); + const int64_t n = ids.shape[0] * table.shape[1]; + if (n == 0) return; + VT_CHECK(table.shape[0] > 0, "rocm embedding: empty table (vocab 0)"); + + hipStream_t s = static_cast(q.handle); + EmbeddingErr* derr = nullptr; + Check(hipMalloc(reinterpret_cast(&derr), sizeof(EmbeddingErr)), + "hipMalloc embedding flag"); + Check(hipMemsetAsync(derr, 0, sizeof(EmbeddingErr), s), "hipMemset embedding flag"); + + if (table.dtype == DType::kF32) { + LaunchEmbeddingIn(s, out, table, ids, derr); + } else { + LaunchEmbeddingIn<__hip_bfloat16>(s, out, table, ids, derr); + } + Check(hipGetLastError(), "embedding launch"); + Check(hipStreamSynchronize(s), "embedding sync"); + + EmbeddingErr herr{}; + Check(hipMemcpy(&herr, derr, sizeof(EmbeddingErr), hipMemcpyDeviceToHost), + "embedding flag D2H"); + Check(hipFree(derr), "hipFree embedding flag"); + + if (herr.status != 0) { + throw std::runtime_error("vt rocm: embedding: id " + std::to_string(herr.id) + + " out of range [0, " + + std::to_string(table.shape[0]) + ")"); + } +} + +} // namespace vt::rocm diff --git a/src/vt/rocm/rocm_gemma4_experts.hip b/src/vt/rocm/rocm_gemma4_experts.hip new file mode 100644 index 000000000..8f274cd69 --- /dev/null +++ b/src/vt/rocm/rocm_gemma4_experts.hip @@ -0,0 +1,111 @@ +// ROCm upload of Gemma-4 MoE expert stacks (BF16 fused or FP8→BF16 dequant). +#include + +#include +#include +#include +#include +#include + +#include "vllm/model_executor/models/gemma4.h" +#include "vllm/model_executor/models/gemma4_moe.h" + +namespace vllm { +namespace { + +void Check(hipError_t e, const char* w) { + if (e != hipSuccess) + throw std::runtime_error(std::string("gemma4 resident: ") + w + ": " + hipGetErrorString(e)); +} + +} // namespace + +size_t UploadGemma4ExpertsResident(std::vector& layers, + int num_gpus) { + if (num_gpus < 1) num_gpus = 1; + int available = 0; + Check(hipGetDeviceCount(&available), "hipGetDeviceCount"); + if (available < 1) return 0; + if (num_gpus > available) num_gpus = available; + + size_t total = 0; + int ok_layers = 0; + int max_layers = 100000; + if (const char* ml = std::getenv("VT_GEMMA4_RESIDENT_MAX_LAYERS")) + max_layers = std::max(0, std::atoi(ml)); + + for (size_t li = 0; li < layers.size(); ++li) { + if (ok_layers >= max_layers) break; + auto& moe = layers[li]; + if (!moe.enabled || moe.experts.Empty()) continue; + auto& ex = moe.experts; + const int64_t E = ex.num_experts; + const int64_t I = ex.intermediate; + const int64_t H = ex.hidden; + const size_t gu_bytes = static_cast(E * 2 * I * H) * 2; + const size_t dn_bytes = static_cast(E * H * I) * 2; + + // Materialize BF16 fused host buffer if FP8 + std::vector gu_host; + std::vector dn_host; + const void* gu_src = nullptr; + const void* dn_src = nullptr; + if (ex.is_fp8) { + gu_host.resize(static_cast(E * 2 * I * H)); + dn_host.resize(static_cast(E * H * I)); + for (int64_t e = 0; e < E; ++e) { + DequantGemma4Fp8ExpertToBf16(ex.fp8[static_cast(e)], I, H, + gu_host.data() + e * 2 * I * H, + dn_host.data() + e * H * I); + } + gu_src = gu_host.data(); + dn_src = dn_host.data(); + } else { + gu_src = ex.gate_up.bytes.data(); + dn_src = ex.down.bytes.data(); + VT_CHECK(ex.gate_up.bytes.size() == gu_bytes && ex.down.bytes.size() == dn_bytes, + "gemma4 resident: bf16 fused size"); + } + + const int dev = static_cast(li % static_cast(num_gpus)); + Check(hipSetDevice(dev), "hipSetDevice"); + void* gu = nullptr; + void* dn = nullptr; + hipError_t e1 = hipMalloc(&gu, gu_bytes); + hipError_t e2 = (e1 == hipSuccess) ? hipMalloc(&dn, dn_bytes) : hipErrorMemoryAllocation; + if (e1 != hipSuccess || e2 != hipSuccess) { + if (gu) (void)hipFree(gu); + if (dn) (void)hipFree(dn); + std::fprintf(stderr, + "gemma4 moe: resident upload stopped at layer %zu on gpu %d (%s)\n", li, + dev, hipGetErrorString(e1 != hipSuccess ? e1 : e2)); + break; + } + Check(hipMemcpy(gu, gu_src, gu_bytes, hipMemcpyHostToDevice), "H2D gu"); + Check(hipMemcpy(dn, dn_src, dn_bytes, hipMemcpyHostToDevice), "H2D dn"); + ex.gate_up_dev = gu; + ex.down_dev = dn; + ex.dev_id = dev; + total += gu_bytes + dn_bytes; + ++ok_layers; + } + Check(hipSetDevice(0), "hipSetDevice0"); + std::fprintf(stderr, "gemma4 moe: resident experts gpus=%d layers=%d %.2f GiB\n", num_gpus, + ok_layers, total / (1024.0 * 1024.0 * 1024.0)); + return total; +} + +size_t UploadGemma4ExpertsResidentForWeights(Gemma4Weights& weights, int num_gpus) { + std::vector flat; + std::vector ptrs; + for (auto& layer : weights.layers) { + if (!layer.moe.enabled) continue; + ptrs.push_back(&layer.moe); + flat.push_back(std::move(layer.moe)); + } + size_t n = UploadGemma4ExpertsResident(flat, num_gpus); + for (size_t i = 0; i < ptrs.size(); ++i) *ptrs[i] = std::move(flat[i]); + return n; +} + +} // namespace vllm diff --git a/src/vt/rocm/rocm_matmul_hipblaslt.hip b/src/vt/rocm/rocm_matmul_hipblaslt.hip new file mode 100644 index 000000000..01c8ddf5a --- /dev/null +++ b/src/vt/rocm/rocm_matmul_hipblaslt.hip @@ -0,0 +1,186 @@ +// ROCm dense GEMM via hipBLAS (hipblasGemmEx) + optional hipBLASLt. +// Primary path: hipblasGemmEx with the standard row-major TN trick used by +// many ROCm LLM stacks. Falls back is not needed — GemmEx is required. +// +// MatmulBT: out[M,N] = a[M,K] @ b[N,K]^T (b = Linear weight row-major) +// Matmul: out[M,N] = a[M,K] @ b[K,N] +#include +#include + +#include +#include +#include +#include +#include + +#include "vt/ops.h" + +namespace vt::rocm { +namespace { + +void CheckHip(hipError_t err, const char* what) { + if (err != hipSuccess) { + throw std::runtime_error(std::string("vt rocm: matmul: ") + what + ": " + + hipGetErrorString(err)); + } +} + +const char* BlasName(hipblasStatus_t st) { + switch (st) { + case HIPBLAS_STATUS_SUCCESS: return "SUCCESS"; + case HIPBLAS_STATUS_NOT_INITIALIZED: return "NOT_INITIALIZED"; + case HIPBLAS_STATUS_ALLOC_FAILED: return "ALLOC_FAILED"; + case HIPBLAS_STATUS_INVALID_VALUE: return "INVALID_VALUE"; + case HIPBLAS_STATUS_MAPPING_ERROR: return "MAPPING_ERROR"; + case HIPBLAS_STATUS_EXECUTION_FAILED: return "EXECUTION_FAILED"; + case HIPBLAS_STATUS_INTERNAL_ERROR: return "INTERNAL_ERROR"; + case HIPBLAS_STATUS_NOT_SUPPORTED: return "NOT_SUPPORTED"; + case HIPBLAS_STATUS_ARCH_MISMATCH: return "ARCH_MISMATCH"; + case HIPBLAS_STATUS_HANDLE_IS_NULLPTR: return "HANDLE_IS_NULLPTR"; + case HIPBLAS_STATUS_INVALID_ENUM: return "INVALID_ENUM"; + default: return "UNKNOWN"; + } +} + +void CheckBlas(hipblasStatus_t st, const char* what) { + if (st != HIPBLAS_STATUS_SUCCESS) { + throw std::runtime_error(std::string("vt rocm: matmul: ") + what + + ": hipblas " + std::to_string(static_cast(st)) + + " (" + BlasName(st) + ")"); + } +} + +struct BlasCtx { + hipblasHandle_t handle = nullptr; +}; + +BlasCtx GetBlas(int device, hipStream_t stream) { + static std::mutex mu; + static std::unordered_map ctxs; + std::lock_guard lock(mu); + auto it = ctxs.find(device); + if (it == ctxs.end()) { + BlasCtx c; + CheckBlas(hipblasCreate(&c.handle), "hipblasCreate"); + ctxs.emplace(device, c); + it = ctxs.find(device); + } + CheckBlas(hipblasSetStream(it->second.handle, stream), "hipblasSetStream"); + return it->second; +} + +std::string ComboName(const Tensor& a, const Tensor& b, const Tensor& out) { + return std::string("(") + Name(a.dtype) + "," + Name(b.dtype) + ")->" + Name(out.dtype); +} + +hipDataType ToBlasType(DType dt) { + switch (dt) { + case DType::kF32: return HIP_R_32F; + case DType::kBF16: return HIP_R_16BF; + case DType::kF16: return HIP_R_16F; + default: + throw std::runtime_error("vt rocm: matmul: unsupported dtype"); + } +} + +} // namespace + +// out[M,N] = a[M,K] @ b[K,N] (row-major contiguous / a may have row stride) +// Col-major BLAS view: C^T = B^T @ A^T +// gemm(OP_T, OP_T, N, M, K, B, ldb=N? wait) +// Standard row-major C = A @ B via: +// gemm(OP_N, OP_N, N, M, K, B_col, N, A_col, K, C_col, N) +// where A_row[M,K] is treated as A_col with ld=K meaning A_col is K x M... +// +// PyTorch/rocBLAS convention for row-major C(m,n) = A(m,k) @ B(k,n): +// hipblasGemm(handle, OP_N, OP_N, n, m, k, +// B, n, A, k, C, n) +// with A,B,C stored row-major — i.e. pass dimensions swapped so BLAS +// writes the transpose layout that coincides with row-major C. +void MatmulKernelRocm(Queue& q, Tensor& out, const Tensor& a, const Tensor& b) { + const bool bf16 = a.dtype == DType::kBF16 && b.dtype == DType::kBF16; + const bool f32 = a.dtype == DType::kF32 && b.dtype == DType::kF32; + if (!bf16 && !f32) { + throw std::runtime_error("vt rocm: matmul: unsupported " + ComboName(a, b, out)); + } + if (out.dtype != DType::kF32 && out.dtype != DType::kBF16) { + throw std::runtime_error("vt rocm: matmul: bad out dtype"); + } + const int64_t M = a.shape[0], K = a.shape[1], N = b.shape[1]; + if (M == 0 || N == 0) return; + hipStream_t s = static_cast(q.handle); + if (K == 0) { + CheckHip(hipMemsetAsync(out.data, 0, out.Bytes(), s), "k0"); + return; + } + // Require contiguous B rows; A may be strided in dim0. + if (b.stride[0] != N) { + throw std::runtime_error("vt rocm: matmul: b must be row-contiguous"); + } + if (a.stride[0] < K) { + throw std::runtime_error("vt rocm: matmul: bad a stride"); + } + + auto ctx = GetBlas(q.device.index, s); + const float alpha = 1.f, beta = 0.f; + const hipDataType at = ToBlasType(a.dtype); + const hipDataType ot = ToBlasType(out.dtype); + + // C_rm = A_rm @ B_rm <=> treat as gemm(N,N) writing C with ld=N, dims n,m,k + // A leading dim = a.stride[0] (row stride in elements) + CheckBlas(hipblasGemmEx(ctx.handle, HIPBLAS_OP_N, HIPBLAS_OP_N, + /*m=*/static_cast(N), /*n=*/static_cast(M), + /*k=*/static_cast(K), &alpha, b.data, at, + /*ldb=*/static_cast(N), a.data, at, + /*lda=*/static_cast(a.stride[0]), &beta, out.data, ot, + /*ldc=*/static_cast(N), HIPBLAS_COMPUTE_32F, + HIPBLAS_GEMM_DEFAULT), + "hipblasGemmEx NN"); +} + +// out[M,N] = a[M,K] @ b[N,K]^T +// Row-major trick: gemm(OP_T, OP_N, N, M, K, B, K, A, a_rs, C, N) +// BLAS: C = op(A)*op(B) with opA=T => A is KxN in col form = row B[N,K] +// opB=N, B is KxM col = row A[M,K] with ld=a_rs +// C is NxM col = row out[M,N] ld=N +void MatmulBTKernelRocm(Queue& q, Tensor& out, const Tensor& a, const Tensor& b) { + const bool bf16 = a.dtype == DType::kBF16 && b.dtype == DType::kBF16; + const bool f32 = a.dtype == DType::kF32 && b.dtype == DType::kF32; + if (!bf16 && !f32) { + throw std::runtime_error("vt rocm: matmul_bt: unsupported " + ComboName(a, b, out)); + } + if (out.dtype != DType::kF32 && out.dtype != DType::kBF16) { + throw std::runtime_error("vt rocm: matmul_bt: bad out dtype"); + } + const int64_t M = a.shape[0], K = a.shape[1], N = b.shape[0]; + if (M == 0 || N == 0) return; + hipStream_t s = static_cast(q.handle); + if (K == 0) { + CheckHip(hipMemsetAsync(out.data, 0, out.Bytes(), s), "bt k0"); + return; + } + if (b.stride[0] != K) { + throw std::runtime_error("vt rocm: matmul_bt: b must be [N,K] row-contiguous"); + } + if (a.stride[0] < K) { + throw std::runtime_error("vt rocm: matmul_bt: bad a stride"); + } + + auto ctx = GetBlas(q.device.index, s); + const float alpha = 1.f, beta = 0.f; + const hipDataType at = ToBlasType(a.dtype); + const hipDataType ot = ToBlasType(out.dtype); + + // C = B^T @ A^T in row-major disguise: + // gemm(OP_T, OP_N, N, M, K, B, ldb=K, A, lda=a_rs, C, ldc=N) + CheckBlas(hipblasGemmEx(ctx.handle, HIPBLAS_OP_T, HIPBLAS_OP_N, + /*m=*/static_cast(N), /*n=*/static_cast(M), + /*k=*/static_cast(K), &alpha, b.data, at, + /*ldb=*/static_cast(K), a.data, at, + /*lda=*/static_cast(a.stride[0]), &beta, out.data, ot, + /*ldc=*/static_cast(N), HIPBLAS_COMPUTE_32F, + HIPBLAS_GEMM_DEFAULT), + "hipblasGemmEx BT"); +} + +} // namespace vt::rocm diff --git a/src/vt/rocm/rocm_ops.hip b/src/vt/rocm/rocm_ops.hip index 206db853b..5667cf00f 100644 --- a/src/vt/rocm/rocm_ops.hip +++ b/src/vt/rocm/rocm_ops.hip @@ -1,41 +1,95 @@ -// ROCm op-table registration (BACKEND-ROCM, W0) — seam 3 of .agents/backends.md. -// -// **UNBUILT.** Never compiled: no AMD GPU on the authoring machine. -// -// This file is the whole of "adding an op to a backend": a static Registrar that -// calls RegisterOp(OpId, DeviceType::kROCM, fn) once per kernel, exactly as -// src/vt/cuda/cuda_ops.cu:3512-3520 and src/vt/vulkan/vulkan_ops.cpp do. There is -// no selector to edit, no model to teach, no dispatch table to extend: an op is -// registered or it is not, and vt::GetOp answers accordingly. -// -// ONE op is registered today (RmsNorm). What happens to the other 105 OpIds -// depends on the board, and both behaviours are correct: -// * UNIFIED memory (Radeon 780M, Strix Halo): the portable reference tier -// (include/vt/op_provider.h:186-224) installs lazily on the first GetOp miss -// and serves them from the CPU kernels, so a model runs end to end, slowly. -// VT_OP_PROVIDER_STATS=1 prints each fallback the first time it is taken — -// that print-out IS the prioritized kernel to-do list for M3. -// * DISCRETE (7900 XTX): the tier is not eligible and GetOp throws. Register -// the op or it does not run. This is the memory-safety gate, not a policy. +// ROCm op-table registration (BACKEND-ROCM). #include "vt/op_provider.h" #include "vt/ops.h" namespace vt::rocm { -// Defined in rocm_rmsnorm.hip. void RmsNormKernelRocm(Queue& q, Tensor& out, const Tensor& x, const Tensor& w, const RmsNormArgs& args, Tensor* residual); +void EmbeddingKernelRocm(Queue& q, Tensor& out, const Tensor& table, const Tensor& ids); +void MulScalarKernelRocm(Queue& q, Tensor& out, const Tensor& x, double scalar); +void SoftCapKernelRocm(Queue& q, Tensor& out, const Tensor& x, double cap); +void AddKernelRocm(Queue& q, Tensor& out, const Tensor& a, const Tensor& b); +void GeluAndMulKernelRocm(Queue& q, Tensor& out, const Tensor& x); +void SiluAndMulKernelRocm(Queue& q, Tensor& out, const Tensor& x); +void MatmulKernelRocm(Queue& q, Tensor& out, const Tensor& a, const Tensor& b); +void MatmulBTKernelRocm(Queue& q, Tensor& out, const Tensor& a, const Tensor& b); +void RopeCosSinCacheKernelRocm(Queue& q, Tensor& cos_sin, const Tensor& pos, + const RopeArgs& args); +void GreedyArgmaxKernelRocm(Queue& q, Tensor& out, const Tensor& logits); +void CastBf16KernelRocm(Queue& q, Tensor& out, const Tensor& in); +void CastF32KernelRocm(Queue& q, Tensor& out, const Tensor& in); +void QkvSplitKernelRocm(Queue& q, Tensor& q_out, Tensor& k_out, Tensor& v_out, + const Tensor& qkv); +void LayerNormKernelRocm(Queue& q, Tensor& out, const Tensor& x, const Tensor* weight, + const Tensor* bias, const LayerNormArgs& args); +void ReluKernelRocm(Queue& q, Tensor& out, const Tensor& x); +void RopeNeoxKernelRocm(Queue& q, Tensor& qs, Tensor& ks, const Tensor& pos, + const RopeArgs& args); +void RopeFromCacheKernelRocm(Queue& q, Tensor& qs, Tensor* ks, const Tensor& positions, + const Tensor& cache, const RopeArgs& args); +void ReshapeAndCacheKernelRocm(Queue& q, const Tensor& k, const Tensor& v, Tensor& k_cache, + Tensor& v_cache, const Tensor& slot_mapping); +void PagedAttentionKernelRocm(Queue& q, Tensor& out, const Tensor& query, const Tensor& k_cache, + const Tensor& v_cache, const Tensor& block_table, + const Tensor& seq_lens, const Tensor& query_start_loc, + const PagedAttentionArgs& args); +void GeluTanhKernelRocm(Queue& q, Tensor& out, const Tensor& x); +void GeluErfKernelRocm(Queue& q, Tensor& out, const Tensor& x); namespace { -// Registers during static init, before main(), like every other backend's op -// table. Harmless on a machine with no AMD GPU: filling the table registers -// nothing about devices, and the kROCM backend never registers there, so no -// kROCM queue can exist to dispatch with. struct Registrar { Registrar() { RegisterOp(OpId::kRmsNorm, DeviceType::kROCM, reinterpret_cast(static_cast(&RmsNormKernelRocm))); + RegisterOp(OpId::kEmbedding, DeviceType::kROCM, + reinterpret_cast(static_cast(&EmbeddingKernelRocm))); + RegisterOp(OpId::kMulScalar, DeviceType::kROCM, + reinterpret_cast(static_cast(&MulScalarKernelRocm))); + RegisterOp(OpId::kSoftCap, DeviceType::kROCM, + reinterpret_cast(static_cast(&SoftCapKernelRocm))); + RegisterOp(OpId::kAdd, DeviceType::kROCM, + reinterpret_cast(static_cast(&AddKernelRocm))); + RegisterOp(OpId::kGeluAndMul, DeviceType::kROCM, + reinterpret_cast(static_cast(&GeluAndMulKernelRocm))); + RegisterOp(OpId::kSiluAndMul, DeviceType::kROCM, + reinterpret_cast(static_cast(&SiluAndMulKernelRocm))); + RegisterOp(OpId::kMatmul, DeviceType::kROCM, + reinterpret_cast(static_cast(&MatmulKernelRocm))); + RegisterOp(OpId::kMatmulBT, DeviceType::kROCM, + reinterpret_cast(static_cast(&MatmulBTKernelRocm))); + RegisterOp(OpId::kRopeCosSinCache, DeviceType::kROCM, + reinterpret_cast( + static_cast(&RopeCosSinCacheKernelRocm))); + RegisterOp(OpId::kGreedyArgmax, DeviceType::kROCM, + reinterpret_cast( + static_cast(&GreedyArgmaxKernelRocm))); + RegisterOp(OpId::kCastBf16, DeviceType::kROCM, + reinterpret_cast(static_cast(&CastBf16KernelRocm))); + RegisterOp(OpId::kCastF32, DeviceType::kROCM, + reinterpret_cast(static_cast(&CastF32KernelRocm))); + RegisterOp(OpId::kQkvSplit, DeviceType::kROCM, + reinterpret_cast(static_cast(&QkvSplitKernelRocm))); + RegisterOp(OpId::kLayerNorm, DeviceType::kROCM, + reinterpret_cast(static_cast(&LayerNormKernelRocm))); + RegisterOp(OpId::kRelu, DeviceType::kROCM, + reinterpret_cast(static_cast(&ReluKernelRocm))); + RegisterOp(OpId::kGeluTanh, DeviceType::kROCM, + reinterpret_cast(static_cast(&GeluTanhKernelRocm))); + RegisterOp(OpId::kGeluErf, DeviceType::kROCM, + reinterpret_cast(static_cast(&GeluErfKernelRocm))); + RegisterOp(OpId::kRopeNeox, DeviceType::kROCM, + reinterpret_cast(static_cast(&RopeNeoxKernelRocm))); + RegisterOp(OpId::kRopeFromCache, DeviceType::kROCM, + reinterpret_cast( + static_cast(&RopeFromCacheKernelRocm))); + RegisterOp(OpId::kReshapeAndCache, DeviceType::kROCM, + reinterpret_cast( + static_cast(&ReshapeAndCacheKernelRocm))); + RegisterOp(OpId::kPagedAttention, DeviceType::kROCM, + reinterpret_cast( + static_cast(&PagedAttentionKernelRocm))); } } registrar; diff --git a/src/vt/rocm/rocm_paged_attn.hip b/src/vt/rocm/rocm_paged_attn.hip new file mode 100644 index 000000000..52bcc1dec --- /dev/null +++ b/src/vt/rocm/rocm_paged_attn.hip @@ -0,0 +1,287 @@ +// ROCm paged attention — port of cpu_paged_attn.cpp algebra + CUDA device-side +// request lookup (cuda_paged_attn.cu). Online softmax; NHD KV via strides. +// Supports bf16/f32 query x bf16/f32 KV (mixed allowed). +#include +#include + +#include +#include +#include +#include + +#include "vt/ops.h" +#include "vt/backend.h" + +#include +#include + +namespace vt::rocm { +namespace { + +constexpr int kBlock = 256; + +inline void Check(hipError_t err, const char* what) { + if (err != hipSuccess) { + throw std::runtime_error(std::string("vt rocm: ") + what + ": " + + hipGetErrorString(err)); + } +} + +// Optional correctness path: run the CPU paged-attn oracle on host copies. +// VT_ROCM_ATTN_CPU_REF=1 — isolates whether custom HIP attn is the decode bug. +bool CpuRefEnabled() { + static const bool on = [] { + const char* e = std::getenv("VT_ROCM_ATTN_CPU_REF"); + return e != nullptr && e[0] == '1' && e[1] == '\0'; + }(); + return on; +} + +void PagedAttnCpuRef(Queue& q, Tensor& out, const Tensor& query, const Tensor& k_cache, + const Tensor& v_cache, const Tensor& block_table, + const Tensor& seq_lens, const Tensor& query_start_loc, + const PagedAttentionArgs& args) { + hipStream_t st = static_cast(q.handle); + Check(hipStreamSynchronize(st), "cpu_ref sync in"); + + // Contiguous host packs (engine KV is a strided unbind view — cannot Bytes()-memcpy). + auto pack4 = [&](const Tensor& t) { + VT_CHECK(t.rank == 4, "cpu_ref pack4"); + const int64_t n0 = t.shape[0], n1 = t.shape[1], n2 = t.shape[2], n3 = t.shape[3]; + const size_t es = vt::SizeOf(t.dtype); + std::vector buf(static_cast(n0 * n1 * n2 * n3) * es); + // element-wise D2H via one big buffer read of parent is hard; use hipMemcpy2D-ish loop + // Simple: copy full bounding box if contiguous innermost + VT_CHECK(t.stride[3] == 1, "cpu_ref need unit inner stride"); + for (int64_t i0 = 0; i0 < n0; ++i0) + for (int64_t i1 = 0; i1 < n1; ++i1) + for (int64_t i2 = 0; i2 < n2; ++i2) { + const int64_t src_off = i0 * t.stride[0] + i1 * t.stride[1] + i2 * t.stride[2]; + const int64_t dst_off = ((i0 * n1 + i1) * n2 + i2) * n3; + Check(hipMemcpy(buf.data() + static_cast(dst_off) * es, + static_cast(t.data) + static_cast(src_off) * es, + static_cast(n3) * es, hipMemcpyDeviceToHost), + "pack4 d2h"); + } + return buf; + }; + auto pack_contig = [&](const Tensor& t) { + std::vector buf(t.Bytes()); + Check(hipMemcpy(buf.data(), t.data, t.Bytes(), hipMemcpyDeviceToHost), "d2h"); + return buf; + }; + + auto h_q = pack_contig(query); + auto h_k = pack4(k_cache); + auto h_v = pack4(v_cache); + auto h_bt = pack_contig(block_table); + auto h_sl = pack_contig(seq_lens); + auto h_qsl = pack_contig(query_start_loc); + std::vector h_out(out.Bytes(), 0); + + vt::Backend& cpu = vt::GetBackend(vt::DeviceType::kCPU); + vt::Queue cq = cpu.CreateQueue(); + const vt::Device cd{vt::DeviceType::kCPU, 0}; + + Tensor tq = Tensor::Contiguous(h_q.data(), query.dtype, cd, + {query.shape[0], query.shape[1], query.shape[2]}); + Tensor tk = Tensor::Contiguous(h_k.data(), k_cache.dtype, cd, + {k_cache.shape[0], k_cache.shape[1], k_cache.shape[2], + k_cache.shape[3]}); + Tensor tv = Tensor::Contiguous(h_v.data(), v_cache.dtype, cd, + {v_cache.shape[0], v_cache.shape[1], v_cache.shape[2], + v_cache.shape[3]}); + Tensor tbt = Tensor::Contiguous(h_bt.data(), block_table.dtype, cd, + {block_table.shape[0], block_table.shape[1]}); + Tensor tsl = Tensor::Contiguous(h_sl.data(), seq_lens.dtype, cd, {seq_lens.shape[0]}); + Tensor tqsl = + Tensor::Contiguous(h_qsl.data(), query_start_loc.dtype, cd, {query_start_loc.shape[0]}); + Tensor to = Tensor::Contiguous(h_out.data(), out.dtype, cd, + {out.shape[0], out.shape[1], out.shape[2]}); + vt::PagedAttention(cq, to, tq, tk, tv, tbt, tsl, tqsl, args); + cpu.DestroyQueue(cq); + + Check(hipMemcpy(out.data, h_out.data(), out.Bytes(), hipMemcpyHostToDevice), "h2d out"); + Check(hipStreamSynchronize(st), "cpu_ref sync out"); +} + +__device__ inline float Ld(const float* p, int64_t i) { return p[i]; } +__device__ inline float Ld(const __hip_bfloat16* p, int64_t i) { + return __bfloat162float(p[i]); +} +__device__ inline void St(float* p, int64_t i, float v) { p[i] = v; } +__device__ inline void St(__hip_bfloat16* p, int64_t i, float v) { + p[i] = __float2bfloat16(v); +} + +__device__ inline float Softcap(float s, float cap) { + return cap > 0.f ? cap * tanhf(s / cap) : s; +} + +// One block per (query_token, q_head). Online softmax matches CUDA. +template +__global__ void PagedAttnOnline(TO* out, const TQ* query, const TKV* k_cache, + const TKV* v_cache, const int32_t* block_table, + const int32_t* seq_lens, const int32_t* query_start_loc, + int64_t num_reqs, int64_t hq, int64_t num_kv_heads, int64_t d, + int64_t block_size, int64_t bt_row, int64_t bt_col, + int64_t kc_blk, int64_t kc_pg, int64_t kc_hd, int64_t vc_blk, + int64_t vc_pg, int64_t vc_hd, float scale, float softcap, + bool causal, int window_left, int window_right) { + const int64_t t = blockIdx.x; + const int64_t h = blockIdx.y; + if (h >= hq) return; + + int64_t r = -1, q0 = 0, q1 = 0; + for (int64_t rr = 0; rr < num_reqs; ++rr) { + const int64_t a = query_start_loc[rr], b = query_start_loc[rr + 1]; + if (t >= a && t < b) { + r = rr; + q0 = a; + q1 = b; + break; + } + } + if (r < 0) return; + + const int64_t query_len = q1 - q0; + const int64_t seqlen = seq_lens[r]; + const int64_t context = seqlen - query_len; + const int64_t p = context + (t - q0); + + int64_t jmin = 0; + if (window_left >= 0) { + jmin = p - window_left; + if (jmin < 0) jmin = 0; + } + int64_t jmax = causal ? p : (seqlen - 1); + if (window_right >= 0) { + const int64_t jr = p + window_right; + if (jr < jmax) jmax = jr; + } + if (jmax > seqlen - 1) jmax = seqlen - 1; + const int64_t g = h / (hq / num_kv_heads); + const int64_t qoff = (t * hq + h) * d; + + // shared: acc[d] + red[blockDim] + extern __shared__ float smem[]; + float* acc = smem; + float* red = smem + d; + __shared__ float s_score, s_m, s_l; + + for (int64_t e = threadIdx.x; e < d; e += blockDim.x) acc[e] = 0.f; + if (threadIdx.x == 0) { + s_m = -INFINITY; + s_l = 0.f; + } + __syncthreads(); + + if (jmax >= jmin) { + for (int64_t j = jmin; j <= jmax; ++j) { + const int64_t blk = block_table[r * bt_row + (j / block_size) * bt_col]; + const int64_t off = j % block_size; + const int64_t kbase = blk * kc_blk + off * kc_pg + g * kc_hd; + float part = 0.f; + for (int64_t e = threadIdx.x; e < d; e += blockDim.x) + part += Ld(query, qoff + e) * Ld(k_cache, kbase + e); + red[threadIdx.x] = part; + __syncthreads(); + for (int stride = (int)blockDim.x / 2; stride > 0; stride >>= 1) { + if ((int)threadIdx.x < stride) red[threadIdx.x] += red[threadIdx.x + stride]; + __syncthreads(); + } + if (threadIdx.x == 0) s_score = Softcap(red[0] * scale, softcap); + __syncthreads(); + + const float s = s_score; + const float m_new = fmaxf(s_m, s); + const float corr = expf(s_m - m_new); + const float pw = expf(s - m_new); + const int64_t vbase = blk * vc_blk + off * vc_pg + g * vc_hd; + for (int64_t e = threadIdx.x; e < d; e += blockDim.x) + acc[e] = acc[e] * corr + pw * Ld(v_cache, vbase + e); + __syncthreads(); + if (threadIdx.x == 0) { + s_l = s_l * corr + pw; + s_m = m_new; + } + __syncthreads(); + } + } + + const float inv = (s_l > 0.f) ? (1.f / s_l) : 0.f; + for (int64_t e = threadIdx.x; e < d; e += blockDim.x) St(out, qoff + e, acc[e] * inv); +} + +} // namespace + +void PagedAttentionKernelRocm(Queue& q, Tensor& out, const Tensor& query, const Tensor& k_cache, + const Tensor& v_cache, const Tensor& block_table, + const Tensor& seq_lens, const Tensor& query_start_loc, + const PagedAttentionArgs& args) { + VT_CHECK(args.kv_cache_dtype == Fp8KVCacheDataType::kAuto, + "rocm paged_attention: fp8 KV not implemented"); + VT_CHECK(args.scale > 0.f, "rocm paged_attention: scale must be > 0"); + if (CpuRefEnabled()) { + PagedAttnCpuRef(q, out, query, k_cache, v_cache, block_table, seq_lens, query_start_loc, + args); + return; + } + const int64_t num_reqs = seq_lens.shape[0]; + const int64_t total_q = query.shape[0]; + const int64_t hq = query.shape[1], d = query.shape[2]; + const int64_t block_size = k_cache.shape[1]; + const int64_t num_kv_heads = k_cache.shape[2]; + VT_CHECK(hq % num_kv_heads == 0, "rocm paged_attention: GQA ratio"); + if (total_q == 0) return; + + hipStream_t st = static_cast(q.handle); + const int window_left = + args.window_size.has_value() ? static_cast(args.window_size->left) : -1; + const int window_right = + args.window_size.has_value() ? static_cast(args.window_size->right) : -1; + const int64_t bt_row = block_table.stride[0], bt_col = block_table.stride[1]; + const int64_t kc_blk = k_cache.stride[0], kc_pg = k_cache.stride[1], kc_hd = k_cache.stride[2]; + const int64_t vc_blk = v_cache.stride[0], vc_pg = v_cache.stride[1], vc_hd = v_cache.stride[2]; + + // threads: reduce over head_dim; need power-of-2 <= 256 and <= d*something + unsigned threads = 64; + while (threads < 256u && threads < static_cast(d)) threads *= 2; + if (threads > 256u) threads = 256; + const size_t smem = sizeof(float) * (static_cast(d) + threads); + dim3 grid(static_cast(total_q), static_cast(hq)); + + auto launch = [&](auto q_tag, auto kv_tag, auto o_tag) { + using TQ = decltype(q_tag); + using TKV = decltype(kv_tag); + using TO = decltype(o_tag); + PagedAttnOnline<<>>( + out.Ptr(), query.Ptr(), k_cache.Ptr(), v_cache.Ptr(), + block_table.Ptr(), seq_lens.Ptr(), query_start_loc.Ptr(), + num_reqs, hq, num_kv_heads, d, block_size, bt_row, bt_col, kc_blk, kc_pg, kc_hd, + vc_blk, vc_pg, vc_hd, args.scale, args.logits_soft_cap, args.causal, window_left, + window_right); + }; + + if (query.dtype == DType::kBF16 && k_cache.dtype == DType::kBF16 && + out.dtype == DType::kBF16) { + launch(__hip_bfloat16{}, __hip_bfloat16{}, __hip_bfloat16{}); + } else if (query.dtype == DType::kF32 && k_cache.dtype == DType::kF32 && + out.dtype == DType::kF32) { + launch(float{}, float{}, float{}); + } else if (query.dtype == DType::kBF16 && k_cache.dtype == DType::kBF16 && + out.dtype == DType::kF32) { + launch(__hip_bfloat16{}, __hip_bfloat16{}, float{}); + } else if (query.dtype == DType::kBF16 && k_cache.dtype == DType::kF32 && + out.dtype == DType::kBF16) { + launch(__hip_bfloat16{}, float{}, __hip_bfloat16{}); + } else if (query.dtype == DType::kF32 && k_cache.dtype == DType::kBF16 && + out.dtype == DType::kF32) { + launch(float{}, __hip_bfloat16{}, float{}); + } else { + VT_CHECK(false, "rocm paged_attention dtype"); + } + Check(hipGetLastError(), "paged_attention launch"); +} + +} // namespace vt::rocm diff --git a/tests/vt/test_backend_cross_device.cpp b/tests/vt/test_backend_cross_device.cpp index d0f486762..ee61e4ab5 100644 --- a/tests/vt/test_backend_cross_device.cpp +++ b/tests/vt/test_backend_cross_device.cpp @@ -494,6 +494,95 @@ TEST_CASE("ReshapeAndCache scatters into the KV cache BIT-EXACTLY") { dev.Free(dsm); dev.DestroyQueue(q); } + + // --- Unbind flash layout: single (blocks,2,bs,H,D) allocation, K/V strided --- + // Matches dense_attn::KvSlice — the layout the engine really feeds. + { + const int64_t within = kBS * kHk * kD; + std::vector combined(static_cast(kBlocks * 2 * within)); + for (int64_t b = 0; b < kBlocks; ++b) + for (int64_t e = 0; e < within; ++e) { + combined[static_cast((b * 2 + 0) * within + e)] = + kc0[static_cast(b * within + e)]; + combined[static_cast((b * 2 + 1) * within + e)] = + vc0[static_cast(b * within + e)]; + } + std::vector ref_comb = combined; + { + vt::Backend& cpu = vt::GetBackend(DeviceType::kCPU); + Queue cq = cpu.CreateQueue(); + const Device cd{DeviceType::kCPU, 0}; + std::vector ck = knew, cv = vnew, cslots_f; + std::vector cslots = slots; + Tensor tk = Tensor::Contiguous(ck.data(), DType::kF32, cd, {kTokens, kHk, kD}); + Tensor tv = Tensor::Contiguous(cv.data(), DType::kF32, cd, {kTokens, kHk, kD}); + Tensor tcomb = + Tensor::Contiguous(ref_comb.data(), DType::kF32, cd, {kBlocks * 2 * within}); + auto slice = [&](int which) { + Tensor t = tcomb; + t.data = static_cast(t.data) + + static_cast(which) * static_cast(within) * sizeof(float); + t.rank = 4; + t.shape[0] = kBlocks; + t.shape[1] = kBS; + t.shape[2] = kHk; + t.shape[3] = kD; + t.stride[0] = 2 * within; + t.stride[1] = kHk * kD; + t.stride[2] = kD; + t.stride[3] = 1; + return t; + }; + Tensor tsm = Tensor::Contiguous(cslots.data(), DType::kI64, cd, {kTokens}); + Tensor tkc = slice(0), tvc = slice(1); + vt::ReshapeAndCache(cq, tk, tv, tkc, tvc, tsm); + cpu.DestroyQueue(cq); + } + + for (DeviceType dt : RegisteredDevices()) { + if (!OpAvailable(vt::OpId::kReshapeAndCache, dt)) continue; + CAPTURE(DeviceName(dt)); + CAPTURE("unbind"); + vt::Backend& dev = vt::GetBackend(dt); + Queue q = dev.CreateQueue(); + const Device d{dt, 0}; + DevBuf dk(dev, q, kTokens * kElems), dv(dev, q, kTokens * kElems), + dcomb(dev, q, kBlocks * 2 * within); + dk.Upload(knew); + dv.Upload(vnew); + dcomb.Upload(combined); + void* dsm = dev.Alloc(kTokens * sizeof(int64_t)); + dev.Copy(q, dsm, slots.data(), kTokens * sizeof(int64_t)); + dev.Synchronize(q); + Tensor tk = Tensor::Contiguous(dk.ptr(), DType::kF32, d, {kTokens, kHk, kD}); + Tensor tv = Tensor::Contiguous(dv.ptr(), DType::kF32, d, {kTokens, kHk, kD}); + Tensor tcomb = + Tensor::Contiguous(dcomb.ptr(), DType::kF32, d, {kBlocks * 2 * within}); + auto slice = [&](int which) { + Tensor t = tcomb; + t.data = static_cast(t.data) + + static_cast(which) * static_cast(within) * sizeof(float); + t.rank = 4; + t.shape[0] = kBlocks; + t.shape[1] = kBS; + t.shape[2] = kHk; + t.shape[3] = kD; + t.stride[0] = 2 * within; + t.stride[1] = kHk * kD; + t.stride[2] = kD; + t.stride[3] = 1; + return t; + }; + Tensor tsm = Tensor::Contiguous(dsm, DType::kI64, d, {kTokens}); + Tensor tkc = slice(0), tvc = slice(1); + vt::ReshapeAndCache(q, tk, tv, tkc, tvc, tsm); + dev.Synchronize(q); + const std::vector got = dcomb.Download(); + CHECK(std::memcmp(ref_comb.data(), got.data(), ref_comb.size() * sizeof(float)) == 0); + dev.Free(dsm); + dev.DestroyQueue(q); + } + } } TEST_CASE("paged attention matches the CPU oracle within NMSE <= 5e-4") { @@ -590,6 +679,71 @@ TEST_CASE("paged attention matches the CPU oracle within NMSE <= 5e-4") { dev.DestroyQueue(q); } } + + // --- DECODE shape: one new query token over a filled cache (Tq=1, seq=kT). + // This is the path multi-token generation hits after prefill; a prefill-only + // test leaves it unexercised. + { + constexpr int64_t kTq = 1; + const std::vector q_dec = RandomVec(kTq * kHq * kD, 701); + const std::vector sl_dec = {static_cast(kT)}; + const std::vector qsl_dec = {0, 1}; + vt::PagedAttentionArgs dargs; + dargs.scale = 0.353553f; + dargs.causal = true; + + vt::Backend& cpu = vt::GetBackend(DeviceType::kCPU); + Queue cq = cpu.CreateQueue(); + const Device cd{DeviceType::kCPU, 0}; + std::vector cq_v = q_dec, ckc = kc, cvc = vc, ref(kTq * kHq * kD); + std::vector cbt = block_table, csl = sl_dec, cqsl = qsl_dec; + { + Tensor tq = Tensor::Contiguous(cq_v.data(), DType::kF32, cd, {kTq, kHq, kD}); + Tensor tkc = Tensor::Contiguous(ckc.data(), DType::kF32, cd, {kBlocks, kBS, kHk, kD}); + Tensor tvc = Tensor::Contiguous(cvc.data(), DType::kF32, cd, {kBlocks, kBS, kHk, kD}); + Tensor tbt = Tensor::Contiguous(cbt.data(), DType::kI32, cd, {1, kBlocks}); + Tensor tsl = Tensor::Contiguous(csl.data(), DType::kI32, cd, {1}); + Tensor tqsl = Tensor::Contiguous(cqsl.data(), DType::kI32, cd, {2}); + Tensor to = Tensor::Contiguous(ref.data(), DType::kF32, cd, {kTq, kHq, kD}); + vt::PagedAttention(cq, to, tq, tkc, tvc, tbt, tsl, tqsl, dargs); + } + cpu.DestroyQueue(cq); + + for (DeviceType dt : RegisteredDevices()) { + if (!OpAvailable(vt::OpId::kPagedAttention, dt)) continue; + CAPTURE(DeviceName(dt)); + CAPTURE("decode"); + vt::Backend& dev = vt::GetBackend(dt); + Queue q = dev.CreateQueue(); + const Device d{dt, 0}; + DevBuf dq(dev, q, kTq * kHq * kD), dkc(dev, q, kBlocks * kBS * kHk * kD), + dvc(dev, q, kBlocks * kBS * kHk * kD), dout(dev, q, kTq * kHq * kD); + dq.Upload(q_dec); + dkc.Upload(kc); + dvc.Upload(vc); + void* dbt = dev.Alloc(kBlocks * sizeof(int32_t)); + void* dsl = dev.Alloc(sizeof(int32_t)); + void* dqsl = dev.Alloc(2 * sizeof(int32_t)); + dev.Copy(q, dbt, block_table.data(), kBlocks * sizeof(int32_t)); + dev.Copy(q, dsl, sl_dec.data(), sizeof(int32_t)); + dev.Copy(q, dqsl, qsl_dec.data(), 2 * sizeof(int32_t)); + dev.Synchronize(q); + Tensor tq = Tensor::Contiguous(dq.ptr(), DType::kF32, d, {kTq, kHq, kD}); + Tensor tkc = Tensor::Contiguous(dkc.ptr(), DType::kF32, d, {kBlocks, kBS, kHk, kD}); + Tensor tvc = Tensor::Contiguous(dvc.ptr(), DType::kF32, d, {kBlocks, kBS, kHk, kD}); + Tensor tbt = Tensor::Contiguous(dbt, DType::kI32, d, {1, kBlocks}); + Tensor tsl = Tensor::Contiguous(dsl, DType::kI32, d, {1}); + Tensor tqsl = Tensor::Contiguous(dqsl, DType::kI32, d, {2}); + Tensor to = Tensor::Contiguous(dout.ptr(), DType::kF32, d, {kTq, kHq, kD}); + vt::PagedAttention(q, to, tq, tkc, tvc, tbt, tsl, tqsl, dargs); + dev.Synchronize(q); + CHECK(Nmse(ref, dout.Download()) <= kNmseTol); + dev.Free(dbt); + dev.Free(dsl); + dev.Free(dqsl); + dev.DestroyQueue(q); + } + } } TEST_CASE("Embedding gather and greedy argmax match the CPU oracle EXACTLY") {