diff --git a/CMakeLists.txt b/CMakeLists.txt index fb49db3da..7e89e81d3 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -840,6 +840,7 @@ add_library(vllm STATIC src/vllm/entrypoints/openai/serving_utils.cpp src/vllm/entrypoints/openai/serving_completion.cpp src/vllm/entrypoints/openai/serving_chat.cpp + src/vllm/entrypoints/openai/request_logger.cpp src/vllm/entrypoints/openai/serving_models.cpp src/vllm/entrypoints/openai/run_batch.cpp src/vllm/entrypoints/openai/tool_parsers/abstract.cpp @@ -912,6 +913,7 @@ add_library(vllm STATIC src/vt/op_provider.cpp src/vt/communicator.cpp src/vt/ops.cpp + src/vt/fused_ops.cpp src/vt/merged_gemm.cpp src/vt/cuda/nvfp4_persistent_cache.cpp src/vt/cpu/cpu_backend.cpp @@ -1145,6 +1147,10 @@ if(VLLM_CPP_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_gemma4_fused_experts.hip + src/vt/rocm/rocm_gemma4_expert_geglu.hip + src/vt/rocm/rocm_fp8_channel_gemv.hip + src/vt/rocm/rocm_moe_router.hip src/vt/rocm/rocm_ops.hip) if(VLLM_CPP_HIP_ARCHITECTURES) set_source_files_properties( @@ -1155,6 +1161,10 @@ if(VLLM_CPP_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_gemma4_fused_experts.hip + src/vt/rocm/rocm_gemma4_expert_geglu.hip + src/vt/rocm/rocm_fp8_channel_gemv.hip + src/vt/rocm/rocm_moe_router.hip src/vt/rocm/rocm_ops.hip PROPERTIES HIP_ARCHITECTURES "${VLLM_CPP_HIP_ARCHITECTURES}") endif() diff --git a/docs/BENCHMARKS.md b/docs/BENCHMARKS.md index 085d118a9..55bb57e1e 100644 --- a/docs/BENCHMARKS.md +++ b/docs/BENCHMARKS.md @@ -373,3 +373,12 @@ built on it rather than keeping the flattering one. Build flags, environment variables, and the full gate list are in [BUILD.md](BUILD.md) and [ENVIRONMENT.md](ENVIRONMENT.md). + +## 2026-08-08 — Gemma4 FP8 stream lab (gfx1201 R9700) + +| Path | Warm tok/s | Notes | +|------|------------|--------| +| vllm-cli Paris HIP=0 stream experts | ~38 | `--repeat` after cold expert fill | +| server `/v1/completions` | ~38 | exclusive | +| server `/v1/chat` thinking off | ~32 | after expert cache | +| llama.cpp Vulkan Q8 tg128 (bar) | ~98 | separate stack | diff --git a/docs/FEATURES.md b/docs/FEATURES.md index 781a9c950..51447d9f6 100644 --- a/docs/FEATURES.md +++ b/docs/FEATURES.md @@ -307,3 +307,5 @@ backends in scope as inventoried rows. Neither changed a capability, so **no mark on this page moved**. An inventoried backend is not a supported one, and the same holds for the 31 architectures inventoried on 2026-08-05. A row's lifecycle state and its support mark are independent: see [STATUS.md](STATUS.md). Parakeet ASR (encoder + CTC/RNN-T/TDT) runs natively on CPU, 4 checkpoints token-exact vs HF. + +| Gemma4 MoE ROCm fused helpers (`vt::fused_ops`) | partial | Portable seam; ROCm fast path; CPU/Vulkan link | diff --git a/docs/STATUS.md b/docs/STATUS.md index 6fbbd7503..4b5530dc2 100644 --- a/docs/STATUS.md +++ b/docs/STATUS.md @@ -2210,3 +2210,11 @@ and outside this repo. **Agent onboarding:** [session](../.agents/specs/session-onboarding.md) + [entry](../.agents/specs/developer-agent-protocol-entrypoint.md) implemented; documentation-only. + +## 2026-08-08 — Gemma4 ROCm fused helpers via portable vt:: seam (#154) + +Model files (`gemma4.cpp`, `gemma4_moe.cpp`) no longer call `vt::rocm::*` directly. +Fused paths go through `include/vt/fused_ops.h` (`vt::RmsNormPlusAdd`, +`DualRmsNormPlusRes`, `GeluMulSeparate`, `MatmulBTAlphaBeta`, `MatmulBTFp8Channel`, +`ExpertGeGLUBf16TopKM1`). ROCm fast path under `VLLM_CPP_HIP`; non-HIP stubs for +peer/pin/resident upload. `check-device-leakage` holds baseline. diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index ed10d4fa4..1a23fac2c 100644 --- a/examples/cli/main.cpp +++ b/examples/cli/main.cpp @@ -8,14 +8,16 @@ // vllm-cli --model --prompt "" // [--tokenizer-config ] [--device auto|cpu|cuda] // [--max-tokens N] [--temperature T] [--top-p P] [--top-k K] -// [--seed S] [--stream] +// [--seed S] [--stream] [--repeat N] // [--gpu-memory-utilization F] [--kv-cache-memory BYTES] // // holds config.json, tokenizer.json and the *.safetensors shards (T0: // safetensors only). Loading a real checkpoint is a GPU/dgx concern; on a CPU // box `--help` / bad-args still work without a model (smoke-tested in CI). +// --repeat N runs N completions after one load (warm bench / decode tok-s). #include "vllm.h" +#include #include #include #include @@ -36,6 +38,7 @@ struct Args { unsigned long long seed = 0; bool have_seed = false; bool stream = false; + int repeat = 1; // load once, complete N times (warm tok/s) std::string speculative_config; // vLLM --speculative-config JSON; "" => off. // --device (ABI v14): "auto" (default probe), "cpu", or "cuda" — the names // of vLLM's DeviceConfig.device this build serves. Mapped to the int the ABI @@ -54,12 +57,13 @@ void Usage(const char* argv0, std::FILE* out) { "usage: %s --model --prompt \"\"\n" " [--tokenizer-config ] [--device auto|cpu|cuda]\n" " [--max-tokens N] [--temperature T] [--top-p P] [--top-k K]\n" - " [--seed S] [--stream]\n" + " [--seed S] [--stream] [--repeat N]\n" " [--gpu-memory-utilization F] [--kv-cache-memory BYTES]\n" " [--speculative-config '']\n" "\n" - "Runs one completion over the vllm.cpp C ABI (libvllm). holds\n" - "config.json, tokenizer.json and the *.safetensors shards.\n", + "Runs completion(s) over the vllm.cpp C ABI (libvllm). holds\n" + "config.json, tokenizer.json and the *.safetensors shards.\n" + "--repeat N: load once, run N blocking completions (for warm tok/s).\n", argv0); } @@ -99,6 +103,9 @@ bool ParseArgs(int argc, char** argv, Args& a, int& exit_code) { a.have_seed = true; } else if (flag == "--stream") { a.stream = true; + } else if (flag == "--repeat") { + a.repeat = std::atoi(NextArg(argc, argv, i)); + if (a.repeat < 1) a.repeat = 1; } else if (flag == "--speculative-config") { a.speculative_config = NextArg(argc, argv, i); } else if (flag == "--gpu-memory-utilization") { @@ -218,6 +225,9 @@ int main(int argc, char** argv) { int rc = 0; if (args.stream) { // ── Streaming: print deltas as they arrive. ────────────────────────────── + if (args.repeat != 1) { + std::fprintf(stderr, "vllm-cli: --repeat with --stream not supported; using 1\n"); + } st = vllm_complete_stream(engine, args.prompt.c_str(), &sp, &StreamPrintCb, nullptr); std::fputc('\n', stdout); @@ -227,20 +237,33 @@ int main(int argc, char** argv) { rc = 1; } } else { - // ── Blocking: run to completion, then print the whole text. ────────────── - vllm_completion out; - st = vllm_complete(engine, args.prompt.c_str(), &sp, &out); - if (st != VLLM_OK) { - std::fprintf(stderr, "vllm-cli: completion failed (status %d): %s\n", - static_cast(st), vllm_last_error()); - rc = 1; - } else { - std::fputs(out.text != nullptr ? out.text : "", stdout); - std::fputc('\n', stdout); + // ── Blocking: load once, optionally repeat for warm tok/s. ─────────────── + for (int r = 0; r < args.repeat; ++r) { + vllm_completion out{}; + const auto t0 = std::chrono::steady_clock::now(); + st = vllm_complete(engine, args.prompt.c_str(), &sp, &out); + const auto t1 = std::chrono::steady_clock::now(); + const double secs = + std::chrono::duration(t1 - t0).count(); + if (st != VLLM_OK) { + std::fprintf(stderr, "vllm-cli: completion failed (status %d): %s\n", + static_cast(st), vllm_last_error()); + rc = 1; + break; + } + if (r == 0 || args.repeat == 1) { + std::fputs(out.text != nullptr ? out.text : "", stdout); + std::fputc('\n', stdout); + } + const int ct = out.completion_tokens; + const double tps = (secs > 0.0 && ct > 0) ? (static_cast(ct) / secs) : 0.0; std::fprintf(stderr, - "vllm-cli: finish_reason=%s prompt_tokens=%d completion_tokens=%d\n", + "vllm-cli: run=%d/%d finish_reason=%s prompt_tokens=%d " + "completion_tokens=%d secs=%.3f tok_s=%.3f\n", + r + 1, args.repeat, out.finish_reason != nullptr ? out.finish_reason : "(none)", - out.prompt_tokens, out.completion_tokens); + out.prompt_tokens, ct, secs, tps); + std::fflush(stderr); vllm_completion_free(&out); } } diff --git a/examples/server/main.cpp b/examples/server/main.cpp index 8580da292..4c9c26edc 100644 --- a/examples/server/main.cpp +++ b/examples/server/main.cpp @@ -62,9 +62,11 @@ #include #include "vllm/entrypoints/openai/api_server.h" #include "vllm/entrypoints/openai/chat_mm.h" +#include "vllm/entrypoints/openai/request_logger.h" #include "vllm/entrypoints/openai/serving_chat.h" #include "vllm/entrypoints/openai/serving_completion.h" #include "vllm/entrypoints/openai/serving_models.h" +#include "vllm/v1/metrics/loggers.h" #include "vllm/entrypoints/openai/reasoning_parsers/detect.h" #include "vllm/entrypoints/openai/tool_parsers/detect.h" #include "vllm/model_executor/model_loader/safetensors_reader.h" @@ -178,6 +180,15 @@ struct Args { // routers only under `if envs.VLLM_SERVER_DEV_MODE` (api_server.py:238). Off by // default → /abort_requests 404s. Enables the /abort_requests production wiring. bool enable_server_dev_mode = false; + bool verbose = false; + // Gemma4 HF/vLLM: --default-chat-template-kwargs enable_thinking (default OFF). + bool enable_thinking = false; + // Request logging (Python vLLM --enable-log-requests parity). Default ON. + bool enable_log_requests = true; + bool enable_log_outputs = false; + int max_log_len = 256; + // Attach Prometheus logger + GET /metrics (default ON for solid Hermes serve). + bool enable_metrics = true; // Scheduling policy: "fcfs" (default), "priority" (mirrors vLLM's // --scheduling-policy / SchedulerConfig.policy), or "lpm" (SGLang's // cache-aware longest-prefix-match admission ordering, ENG-SGLANG-BEHAVIOR-FLAG; @@ -225,6 +236,11 @@ struct Args { " [--enable-force-include-usage]\n" " [--enable-tokenizer-info-endpoint]\n" " [--enable-server-dev-mode]\n" + " [--verbose]\n" + " [--enable-thinking|--no-enable-thinking]\n" + " [--enable-log-requests|--disable-log-requests]\n" + " [--enable-log-outputs] [--max-log-len N]\n" + " [--enable-metrics|--disable-metrics]\n" " [--[no-]enable-prefix-caching]\n" " [--[no-]enable-radix-attention]\n" " [--scheduling-policy fcfs|priority|lpm]\n" @@ -319,6 +335,24 @@ Args ParseArgs(int argc, char** argv) { a.video_dequant_bf16 = true; } else if (flag == "--enable-server-dev-mode") { a.enable_server_dev_mode = true; + } else if (flag == "--verbose" || flag == "-v") { + a.verbose = true; + } else if (flag == "--enable-thinking") { + a.enable_thinking = true; + } else if (flag == "--no-enable-thinking") { + a.enable_thinking = false; + } else if (flag == "--enable-log-requests") { + a.enable_log_requests = true; + } else if (flag == "--disable-log-requests") { + a.enable_log_requests = false; + } else if (flag == "--enable-log-outputs") { + a.enable_log_outputs = true; + } else if (flag == "--max-log-len") { + a.max_log_len = std::stoi(NextArg(argc, argv, i, argv[0])); + } else if (flag == "--enable-metrics") { + a.enable_metrics = true; + } else if (flag == "--disable-metrics") { + a.enable_metrics = false; } else if (flag == "--enable-prefix-caching" || flag == "--no-enable-prefix-caching" || flag == "--enable-radix-attention" || @@ -416,6 +450,26 @@ Args ParseArgs(int argc, char** argv) { int main(int argc, char** argv) { try { const Args args = ParseArgs(argc, argv); + if (args.verbose) { + setenv("VT_SERVER_VERBOSE", "1", /*overwrite=*/1); + std::cerr << "server: verbose stage logging enabled (debug_stages)\n"; + } + { + vllm::entrypoints::openai::RequestLogConfig log_cfg; + log_cfg.enable_log_requests = args.enable_log_requests; + log_cfg.enable_log_outputs = args.enable_log_outputs || args.verbose; + log_cfg.max_log_len = args.max_log_len; + log_cfg.debug_stages = args.verbose || + (std::getenv("VT_SERVER_VERBOSE") && + std::getenv("VT_SERVER_VERBOSE")[0] == '1'); + vllm::entrypoints::openai::ConfigureRequestLogger(log_cfg); + std::cerr << "server: request logging " + << (log_cfg.enable_log_requests ? "ON" : "OFF") + << " outputs=" << (log_cfg.enable_log_outputs ? "ON" : "OFF") + << " max_log_len=" << log_cfg.max_log_len + << " debug_stages=" << (log_cfg.debug_stages ? "ON" : "OFF") + << "\n"; + } const fs::path dir(args.model_dir); const std::string config_path = (dir / "config.json").string(); @@ -666,9 +720,13 @@ int main(int argc, char** argv) { const std::string eos = tokenizer.EosId() >= 0 ? tokenizer.Decode({tokenizer.EosId()}) : ""; chat_prompt_fn = - vllm::entrypoints::MakeChatTemplatePromptFn(chat_template, bos, eos); - std::cerr << "server: using chat template from " << tokenizer_config_path - << "\n"; + vllm::entrypoints::MakeChatTemplatePromptFn( + chat_template, bos, eos, args.enable_thinking); + std::cerr << "server: using chat template (" << chat_template.size() + << " chars) from " << tokenizer_config_path + << " or sibling chat_template.jinja" + << " enable_thinking=" + << (args.enable_thinking ? "true" : "false") << "\n"; } catch (const std::exception& e) { std::cerr << "server: no chat template (" << e.what() << "); falling back to the simple role-join prompt\n"; @@ -855,6 +913,17 @@ int main(int argc, char** argv) { : "") << "\n"; + // Prometheus /metrics (Python vLLM always-on family names). + std::unique_ptr prom_logger; + if (args.enable_metrics) { + prom_logger = std::make_unique( + served_model_name, loaded->max_model_len(), /*engine_index=*/0); + // Sync engine path records on step(); async may under-report until fully wired. + loaded->engine().set_stat_logger(prom_logger.get()); + server.set_metrics_logger(prom_logger.get()); + std::cerr << "server: GET /metrics enabled (PrometheusStatLogger)\n"; + } + std::cerr << "server: listening on http://" << args.host << ":" << args.port << " (model '" << served_model_name << "', HTTP worker pool "; if (server.http_worker_count() == 0) { diff --git a/include/vllm/entrypoints/chat_template.h b/include/vllm/entrypoints/chat_template.h index 8fe7eb3ea..8c0daae6d 100644 --- a/include/vllm/entrypoints/chat_template.h +++ b/include/vllm/entrypoints/chat_template.h @@ -63,19 +63,22 @@ class ChatTemplateError : public std::runtime_error { // branch. Empty => the `tools` variable is an empty list // (falsy). The `tojson` filter is a minja builtin. // Throws ChatTemplateError on any parse or evaluation error. +// chat_template_kwargs: optional Jinja variables (vLLM +// --default-chat-template-kwargs). Supported keys today: enable_thinking (bool). std::string apply_chat_template( const std::string& template_str, const std::vector& messages, bool add_generation_prompt, const std::string& bos_token = "", const std::string& eos_token = "", - const std::vector& tools = {}); + const std::vector& tools = {}, + bool enable_thinking = false); -// Adapt a chat template to Task 2's ChatPromptFn seam (serving_chat.h). The -// returned callable renders `template_str` for the messages + generation flag it -// is handed, so an OpenAIServingChat constructed with it applies the real chat -// template instead of DefaultChatPromptFallback. +// Adapt a chat template to Task 2's ChatPromptFn seam (serving_chat.h). +// enable_thinking defaults false for agent/Hermes latency (Gemma4 empty thought +// block when false — HF/vLLM recipe parity). openai::ChatPromptFn MakeChatTemplatePromptFn(std::string template_str, std::string bos_token = "", - std::string eos_token = ""); + std::string eos_token = "", + bool enable_thinking = false); // Load the `chat_template` string out of a tokenizer_config.json file. Handles // both the plain-string form and the list-of-{name,template} form (picks the diff --git a/include/vllm/entrypoints/openai/request_logger.h b/include/vllm/entrypoints/openai/request_logger.h new file mode 100644 index 000000000..b32f69935 --- /dev/null +++ b/include/vllm/entrypoints/openai/request_logger.h @@ -0,0 +1,49 @@ +// OpenAI serve request logging — shaped like Python vLLM --enable-log-requests. +// Ported concepts from: vllm/entrypoints/logger.py + api_server access logs. +#pragma once + +#include +#include + +namespace vllm::entrypoints::openai { + +struct RequestLogConfig { + // Mirrors --enable-log-requests / --disable-log-requests (default ON for serve). + bool enable_log_requests = true; + // Mirrors --enable-log-outputs (requires enable_log_requests). + bool enable_log_outputs = false; + // Mirrors --max-log-len (chars of prompt/output preview). + int max_log_len = 256; + // Lab deep stages (former VT_SERVER_VERBOSE chat-dbg). + bool debug_stages = false; +}; + +// Process-wide config (set once at server startup). +void ConfigureRequestLogger(const RequestLogConfig& cfg); +const RequestLogConfig& GetRequestLogConfig(); + +// Truncate for log lines (max_log_len); escapes newlines. +std::string LogPreview(const std::string& s, int max_len); + +// HTTP ingress (api_server). +void LogHttpIngress(const char* method, const char* path, size_t body_bytes); + +// After chat/completions parse + template. +void LogRequestReceived(const std::string& request_id, const std::string& endpoint, + const std::string& model, bool stream, int max_tokens, + int n_messages, int n_tools, size_t prompt_chars, + const std::string& prompt, const std::string& roles_summary); + +// Mid-flight stages (debug_stages only) — heartbeats during prefill/decode. +void LogRequestStage(const std::string& request_id, const std::string& stage); + +// Completion of a request. +void LogRequestFinished(const std::string& request_id, int prompt_tokens, + int completion_tokens, const std::string& finish_reason, + double elapsed_sec, const std::string& output_text); + +// Errors. +void LogRequestError(const std::string& request_id, const std::string& endpoint, + const std::string& what); + +} // namespace vllm::entrypoints::openai diff --git a/include/vllm/model_executor/model_loader/nvfp4_dequant.h b/include/vllm/model_executor/model_loader/nvfp4_dequant.h index 680c26e48..2c73b3bc2 100644 --- a/include/vllm/model_executor/model_loader/nvfp4_dequant.h +++ b/include/vllm/model_executor/model_loader/nvfp4_dequant.h @@ -84,4 +84,8 @@ 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); +// Nesting guards: expert-level parallel prefetch serializes row-parallel dequant. +void Fp8DequantBeginOuterParallel(); +void Fp8DequantEndOuterParallel(); + } // namespace vllm diff --git a/include/vllm/model_executor/models/gemma4_moe.h b/include/vllm/model_executor/models/gemma4_moe.h index f288a4a56..a04864094 100644 --- a/include/vllm/model_executor/models/gemma4_moe.h +++ b/include/vllm/model_executor/models/gemma4_moe.h @@ -25,9 +25,15 @@ struct Gemma4Fp8ExpertMats { // Lazy host BF16 cache after first dequant (decode reuse). mutable std::vector cached_gu; // [2I,H] mutable std::vector cached_dn; // [H,I] + mutable bool host_pinned = false; // 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 + // Native FP8 on device (VT_GEMMA4_FP8_NATIVE=1): half weight bandwidth vs BF16. + mutable void* dev_fp8_gu = nullptr; // u8 [2I,H] gate|up + mutable void* dev_fp8_dn = nullptr; // u8 [H,I] + mutable void* dev_s_gu = nullptr; // bf16 [2I] + mutable void* dev_s_dn = nullptr; // bf16 [H] }; struct Gemma4FusedExperts { @@ -74,8 +80,30 @@ size_t UploadGemma4ExpertsResident(std::vector& layers, int num_gpus); size_t UploadGemma4ExpertsResidentForWeights(Gemma4Weights& weights, int num_gpus); +// Peer-copy one resident expert (fused stacks on src_dev) into dst buffers on +// compute_dev. Returns false if peer path unavailable. +bool PeerCopyGemma4ExpertSlice(int src_dev, const void* gate_up_base, + const void* down_base, int expert_id, int64_t I, + int64_t H, int compute_dev, void* gate_up_dst, + void* down_dst); + +// hipHostRegister BF16 expert cache for faster H2D (no-op if already pinned). +void PinGemma4Fp8ExpertHostCache(const Gemma4Fp8ExpertMats& ex); + // Dequant one FP8 expert into host BF16 gate_up[2I,H] and down[H,I] (caller-owned). +// Fills permanent host cache (decode path). Prefer Ephemeral for bulk upload. void DequantGemma4Fp8ExpertToBf16(const Gemma4Fp8ExpertMats& ex, int64_t I, int64_t H, uint16_t* gate_up_out, uint16_t* down_out); +// Same dequant without retaining permanent host BF16 cache (resident upload). +void DequantGemma4Fp8ExpertToBf16Ephemeral(const Gemma4Fp8ExpertMats& ex, int64_t I, + int64_t H, uint16_t* gate_up_out, + uint16_t* down_out); + +// Fused top-k ExpertGeGLU (T=1). Opt-in VT_GEMMA4_FUSED_EXPERTS=1. +// Returns false if disabled/unsupported — caller uses serial hipBLAS path. +bool RunGemma4FusedTopkExpertGeGLU(vt::Queue& q, void* ysum, const void* x, + const uint16_t* const* gu_ptrs, + const uint16_t* const* dn_ptrs, const float* wts, int G, + int64_t I, int64_t H); } // namespace vllm diff --git a/include/vt/fused_ops.h b/include/vt/fused_ops.h new file mode 100644 index 000000000..acc29c167 --- /dev/null +++ b/include/vt/fused_ops.h @@ -0,0 +1,33 @@ +// Portable fused helpers used by Gemma4 (and reusable elsewhere). +// Model code MUST call these — never vt::rocm::* — so CPU/CUDA/Vulkan link. +#pragma once + +#include + +#include "vt/device.h" +#include "vt/dtype.h" +#include "vt/ops.h" +#include "vt/tensor.h" + +namespace vt { + +void RmsNormPlusAdd(Queue& q, Tensor& out, const Tensor& x, const Tensor& w, + const Tensor& addend, const RmsNormArgs& args); + +void DualRmsNormPlusRes(Queue& q, Tensor& out, const Tensor& x1, const Tensor& w1, + const Tensor& x2, const Tensor& w2, const Tensor& w3, + const Tensor& residual, const RmsNormArgs& args); + +void GeluMulSeparate(Queue& q, void* out, const void* gate, const void* up, int64_t n, + DType dtype); + +void MatmulBTAlphaBeta(Queue& q, void* out, const void* a, const void* b, int M, int N, int K, + float alpha, float beta, DType dtype); + +void MatmulBTFp8Channel(Queue& q, void* out, const void* a, const void* b_fp8, + const void* scale_bf16, int M, int N, int K, float alpha, float beta); + +bool ExpertGeGLUBf16TopKM1(Queue& q, void* ysum, const void* x, const void* const* w_gu, + const void* const* w_dn, const float* wts, int G, int I, int H); + +} // namespace vt diff --git a/include/vt/rocm/rocm_gelu_mul_sep.h b/include/vt/rocm/rocm_gelu_mul_sep.h new file mode 100644 index 000000000..f6ea4e7a3 --- /dev/null +++ b/include/vt/rocm/rocm_gelu_mul_sep.h @@ -0,0 +1,9 @@ +#pragma once +#include "vt/device.h" +#include "vt/dtype.h" + +namespace vt::rocm { +// out[i] = gelu_tanh(gate[i]) * up[i] for i in [0,n) +void GeluMulSeparateRocm(Queue& q, void* out, const void* gate, const void* up, int64_t n, + DType dtype); +} // namespace vt::rocm diff --git a/include/vt/rocm/rocm_gemma4_expert_geglu.h b/include/vt/rocm/rocm_gemma4_expert_geglu.h new file mode 100644 index 000000000..666a28535 --- /dev/null +++ b/include/vt/rocm/rocm_gemma4_expert_geglu.h @@ -0,0 +1,17 @@ +// Custom RDNA4 Gemma-4 Expert GeGLU (decode T=1). +#pragma once + +#include "vt/device.h" + +namespace vt::rocm { + +// y[H] = beta*y + alpha * Down(Gelu(Gate(x))*Up(x)) +// w_gu [2I,H] bf16, w_dn [H,I] bf16, x[H] bf16 +bool ExpertGeGLUBf16M1Rocm(Queue& q, void* y, const void* x, const void* w_gu, const void* w_dn, + int I, int H, float alpha, float beta); + +// Top-k sequential mix into ysum. +bool ExpertGeGLUBf16TopKM1Rocm(Queue& q, void* ysum, const void* x, const void* const* w_gu, + const void* const* w_dn, const float* wts, int G, int I, int H); + +} // namespace vt::rocm diff --git a/include/vt/rocm/rocm_matmul_batch.h b/include/vt/rocm/rocm_matmul_batch.h new file mode 100644 index 000000000..f0fd34131 --- /dev/null +++ b/include/vt/rocm/rocm_matmul_batch.h @@ -0,0 +1,36 @@ +// ROCm: batched MatmulBT helpers for MoE top-k fuse. +#pragma once + +#include "vt/device.h" +#include "vt/dtype.h" + +namespace vt::rocm { + +void MatmulBTStridedBatchKernelRocm(Queue& q, void* out, const void* a, const void* b, + int batch, int M, int N, int K, DType dtype); + +void MatmulBTStridedBatchFullKernelRocm(Queue& q, void* out, const void* a, const void* b, + int batch, int M, int N, int K, long long strideA, + DType dtype); + +// Pointer-array batch (no gather): out_ptrs[g] = a @ b_ptrs[g]^T +void MatmulBTPointerBatchKernelRocm(Queue& q, void** out_ptrs, const void* a, + void** b_ptrs, int batch, int M, int N, int K, + DType dtype); + +// Both A and B per-batch: out_ptrs[g] = a_ptrs[g] @ b_ptrs[g]^T +void MatmulBTPointerBatchABKernelRocm(Queue& q, void** out_ptrs, void** a_ptrs, + void** b_ptrs, int batch, int M, int N, int K, + DType dtype); + +// out = alpha * (a @ b^T) + beta * out — MoE expert mix fuse +// a [M,K], b [N,K], out [M,N], contiguous rows +void MatmulBTAlphaBetaRocm(Queue& q, void* out, const void* a, const void* b, int M, int N, + int K, float alpha, float beta, DType dtype); + +// M=1 BF16 act × FP8 weight [N,K] with BF16 channel scale[N] +void MatmulBTFp8ChannelRocm(Queue& q, void* out, const void* a, const void* b_fp8, + const void* scale_bf16, int M, int N, int K, float alpha, + float beta); + +} // namespace vt::rocm diff --git a/include/vt/rocm/rocm_rmsnorm_plus_add.h b/include/vt/rocm/rocm_rmsnorm_plus_add.h new file mode 100644 index 000000000..ccb27592b --- /dev/null +++ b/include/vt/rocm/rocm_rmsnorm_plus_add.h @@ -0,0 +1,16 @@ +#pragma once +#include "vt/device.h" +#include "vt/dtype.h" +#include "vt/ops.h" +#include "vt/tensor.h" + +namespace vt::rocm { +// out = rmsnorm(x, w) + addend (Gemma-4 residual join) +void RmsNormPlusAddRocm(Queue& q, Tensor& out, const Tensor& x, const Tensor& w, + const Tensor& addend, const RmsNormArgs& args); + +// out = rmsnorm(rmsnorm(x1,w1)+rmsnorm(x2,w2), w3) + residual +void DualRmsNormPlusResRocm(Queue& q, Tensor& out, const Tensor& x1, const Tensor& w1, + const Tensor& x2, const Tensor& w2, const Tensor& w3, + const Tensor& residual, const RmsNormArgs& args); +} // namespace vt::rocm diff --git a/src/vllm/entrypoints/chat_template.cpp b/src/vllm/entrypoints/chat_template.cpp index 7d7b89e77..c565a37fc 100644 --- a/src/vllm/entrypoints/chat_template.cpp +++ b/src/vllm/entrypoints/chat_template.cpp @@ -109,22 +109,9 @@ std::string apply_chat_template( const std::string& template_str, const std::vector& messages, bool add_generation_prompt, const std::string& bos_token, const std::string& eos_token, - const std::vector& tools) { + const std::vector& tools, + bool enable_thinking) { try { - // Render the template LITERALLY, exactly as transformers' - // `apply_chat_template` does: parse the raw Jinja source and render it with - // transformers' whitespace policy (trim_blocks / lstrip_blocks, no trailing - // newline). We deliberately use minja's low-level engine rather than its - // high-level `chat_template` wrapper: that wrapper runs a heuristic - // capability probe (6+ speculative renders per construction, plus stderr - // diagnostics) and a "polyfill" pass that rewrites the message list (merging - // the system role into a user turn, injecting a synthetic tools system - // prompt, ...). None of that is transformers behavior, and the probe - // misfires on templates that do not echo message content verbatim. The - // low-level path is the faithful, quiet, per-request-cheap equivalent. - // - // Parser::parse throws on a syntax error; render() throws on an evaluation - // error. Both surface below as ChatTemplateError. std::shared_ptr root = minja::Parser::parse( template_str, minja::Options{/*trim_blocks=*/true, /*lstrip_blocks=*/true, @@ -133,17 +120,14 @@ std::string apply_chat_template( nlohmann::ordered_json top = nlohmann::ordered_json::object(); top["messages"] = BuildMessages(messages); top["add_generation_prompt"] = add_generation_prompt; - // Context::make's default parent is minja::Context::builtins(), which - // provides the standard Jinja filters/functions/tests (tojson, upper, map, - // selectattr, is-tests, ...) the real templates rely on. + // vLLM/HF: enable_thinking controls Gemma4 CoT channel (default false). + top["enable_thinking"] = enable_thinking; std::shared_ptr context = minja::Context::make(minja::Value(top)); context->set("bos_token", minja::Value(bos_token)); context->set("eos_token", minja::Value(eos_token)); - // An empty tools array is falsy in Jinja (`{% if tools %}` skips), matching - // transformers passing tools=None. + context->set("enable_thinking", minja::Value(enable_thinking)); context->set("tools", minja::Value(BuildTools(tools))); - // Some templates (e.g. Llama 3.x) call strftime_now(fmt) for the date line. const auto now = std::chrono::system_clock::now(); context->set( "strftime_now", @@ -174,14 +158,15 @@ std::string apply_chat_template( openai::ChatPromptFn MakeChatTemplatePromptFn(std::string template_str, std::string bos_token, - std::string eos_token) { + std::string eos_token, + bool enable_thinking) { return [tmpl = std::move(template_str), bos = std::move(bos_token), - eos = std::move(eos_token)]( + eos = std::move(eos_token), enable_thinking]( const std::vector& messages, bool add_generation_prompt, const std::vector& tools) { return apply_chat_template(tmpl, messages, add_generation_prompt, bos, eos, - tools); + tools, enable_thinking); }; } @@ -200,27 +185,43 @@ std::string LoadChatTemplateFromConfig( e.what()); } auto it = doc.find("chat_template"); - if (it == doc.end() || it->is_null()) { - throw ChatTemplateError("tokenizer_config.json has no 'chat_template': " + - tokenizer_config_path); - } - if (it->is_string()) return it->get(); - // List-of-{name,template} form: pick "default", else the first. - if (it->is_array()) { - const nlohmann::json* chosen = nullptr; - for (const auto& entry : *it) { - if (entry.is_object() && entry.value("name", std::string()) == "default") { - chosen = &entry; - break; + if (it != doc.end() && !it->is_null()) { + if (it->is_string()) return it->get(); + // List-of-{name,template} form: pick "default", else the first. + if (it->is_array()) { + const nlohmann::json* chosen = nullptr; + for (const auto& entry : *it) { + if (entry.is_object() && entry.value("name", std::string()) == "default") { + chosen = &entry; + break; + } + } + if (!chosen && !it->empty()) chosen = &it->front(); + if (chosen && chosen->contains("template") && + (*chosen)["template"].is_string()) { + return (*chosen)["template"].get(); } } - if (!chosen && !it->empty()) chosen = &it->front(); - if (chosen && chosen->contains("template") && - (*chosen)["template"].is_string()) { - return (*chosen)["template"].get(); - } + throw ChatTemplateError("unrecognized 'chat_template' shape in " + + tokenizer_config_path); + } + + // Sibling chat_template.jinja (HF layout for Gemma4 / many multimodal models). + std::string dir = tokenizer_config_path; + const auto slash = dir.find_last_of("/\\"); + if (slash != std::string::npos) dir.resize(slash + 1); + else dir.clear(); + const std::string jinja_path = dir + "chat_template.jinja"; + std::ifstream jf(jinja_path, std::ios::binary); + if (jf) { + std::ostringstream ss; + ss << jf.rdbuf(); + std::string tmpl = ss.str(); + if (!tmpl.empty()) return tmpl; } - throw ChatTemplateError("unrecognized 'chat_template' shape in " + + + throw ChatTemplateError("tokenizer_config.json has no 'chat_template' and no " + "sibling chat_template.jinja: " + tokenizer_config_path); } diff --git a/src/vllm/entrypoints/openai/api_server.cpp b/src/vllm/entrypoints/openai/api_server.cpp index 1f692dbe2..6cadaa559 100644 --- a/src/vllm/entrypoints/openai/api_server.cpp +++ b/src/vllm/entrypoints/openai/api_server.cpp @@ -4,6 +4,7 @@ #include "vllm/entrypoints/openai/api_server.h" #include +#include #include #include #include @@ -21,6 +22,7 @@ #include #include "vllm/entrypoints/openai/protocol.h" +#include "vllm/entrypoints/openai/request_logger.h" #include "vllm/tokenizer/tokenizer.h" #include "vllm/v1/engine/async_llm.h" #include "vllm/v1/metrics/loggers.h" @@ -209,6 +211,9 @@ ApiServer::DispatchResult ApiServer::handle_chat_completions( "The model does not support Chat Completions API " "(transcription-only server)"); } + { + LogHttpIngress("POST", "/v1/chat/completions", request_body.size()); + } // chat_completion/api_router.py:53 (create_chat_completion). nlohmann::json body; try { @@ -239,6 +244,7 @@ ApiServer::DispatchResult ApiServer::handle_chat_completions( } catch (const std::exception& e) { std::cerr << "api-server: 500 endpoint=/v1/chat/completions model=" << request.model.value_or("") << " what=" << e.what() << "\n"; + LogRequestError("", "/v1/chat/completions", e.what()); return MakeError(500, "InternalServerError", e.what()); } diff --git a/src/vllm/entrypoints/openai/request_logger.cpp b/src/vllm/entrypoints/openai/request_logger.cpp new file mode 100644 index 000000000..264b60c1b --- /dev/null +++ b/src/vllm/entrypoints/openai/request_logger.cpp @@ -0,0 +1,119 @@ +#include "vllm/entrypoints/openai/request_logger.h" + +#include +#include +#include +#include +#include +#include + +namespace vllm::entrypoints::openai { +namespace { + +RequestLogConfig g_cfg{}; +std::mutex g_mu; +std::chrono::steady_clock::time_point g_t0 = std::chrono::steady_clock::now(); + +int64_t MsSinceStart() { + return std::chrono::duration_cast( + std::chrono::steady_clock::now() - g_t0) + .count(); +} + +void Emit(const std::string& line) { + std::lock_guard lock(g_mu); + // Upstream-ish prefix: INFO level one-liners on stderr (no python logging stack). + std::cerr << "INFO " << line << "\n"; + std::cerr.flush(); +} + +} // namespace + +void ConfigureRequestLogger(const RequestLogConfig& cfg) { + g_cfg = cfg; + if (g_cfg.enable_log_outputs && !g_cfg.enable_log_requests) { + g_cfg.enable_log_outputs = false; + } + if (g_cfg.max_log_len < 16) g_cfg.max_log_len = 16; +} + +const RequestLogConfig& GetRequestLogConfig() { return g_cfg; } + +std::string LogPreview(const std::string& s, int max_len) { + std::string out; + out.reserve(static_cast(max_len) + 8); + const size_t n = std::min(s.size(), static_cast(max_len)); + for (size_t i = 0; i < n; ++i) { + const unsigned char c = static_cast(s[i]); + if (c == '\n') + out += "\\n"; + else if (c == '\r') + out += "\\r"; + else if (c == '\t') + out += "\\t"; + else if (c < 32 || c == 127) + out += '?'; + else + out += static_cast(c); + } + if (s.size() > n) out += "..."; + return out; +} + +void LogHttpIngress(const char* method, const char* path, size_t body_bytes) { + if (!g_cfg.enable_log_requests && !g_cfg.debug_stages) return; + std::ostringstream os; + os << "api: " << method << " " << path << " body_bytes=" << body_bytes + << " t+" << MsSinceStart() << "ms"; + Emit(os.str()); +} + +void LogRequestReceived(const std::string& request_id, const std::string& endpoint, + const std::string& model, bool stream, int max_tokens, + int n_messages, int n_tools, size_t prompt_chars, + const std::string& prompt, const std::string& roles_summary) { + if (!g_cfg.enable_log_requests) return; + std::ostringstream os; + os << "Received request " << request_id << " endpoint=" << endpoint + << " model=" << model << " stream=" << (stream ? "1" : "0") + << " max_tokens=" << max_tokens << " msgs=" << n_messages + << " tools=" << n_tools << " prompt_chars=" << prompt_chars; + if (!roles_summary.empty()) os << " roles=" << roles_summary; + os << " prompt: '" << LogPreview(prompt, g_cfg.max_log_len) << "'"; + Emit(os.str()); +} + +void LogRequestStage(const std::string& request_id, const std::string& stage) { + if (!g_cfg.debug_stages) return; + std::ostringstream os; + os << "chat-dbg t+" << MsSinceStart() << "ms id=" << request_id << " " << stage; + Emit(os.str()); +} + +void LogRequestFinished(const std::string& request_id, int prompt_tokens, + int completion_tokens, const std::string& finish_reason, + double elapsed_sec, const std::string& output_text) { + if (!g_cfg.enable_log_requests) return; + std::ostringstream os; + os << "Finished request " << request_id << " prompt_tokens=" << prompt_tokens + << " completion_tokens=" << completion_tokens + << " total_tokens=" << (prompt_tokens + completion_tokens) + << " finish_reason=" << finish_reason << " elapsed_s=" << elapsed_sec; + if (completion_tokens > 0 && elapsed_sec > 0.001) { + os << " gen_tok_s=" << (static_cast(completion_tokens) / elapsed_sec); + } + if (g_cfg.enable_log_outputs && !output_text.empty()) { + os << " output: '" << LogPreview(output_text, g_cfg.max_log_len) << "'"; + } + Emit(os.str()); +} + +void LogRequestError(const std::string& request_id, const std::string& endpoint, + const std::string& what) { + std::ostringstream os; + os << "ERROR request " << (request_id.empty() ? "-" : request_id) + << " endpoint=" << endpoint << " what=" << what; + Emit(os.str()); +} + +} // namespace vllm::entrypoints::openai diff --git a/src/vllm/entrypoints/openai/serving_chat.cpp b/src/vllm/entrypoints/openai/serving_chat.cpp index b704bf76f..b1a3a260b 100644 --- a/src/vllm/entrypoints/openai/serving_chat.cpp +++ b/src/vllm/entrypoints/openai/serving_chat.cpp @@ -2,8 +2,12 @@ // See serving_chat.h for scope, the chat-prompt seam and deferrals. #include "vllm/entrypoints/openai/serving_chat.h" +#include #include +#include +#include #include +#include #include #include #include @@ -13,6 +17,7 @@ #include "vllm/entrypoints/beam_search.h" #include "vllm/entrypoints/openai/chat_mm.h" // HasMultiModalParts (mm seam gate) +#include "vllm/entrypoints/openai/request_logger.h" #include "vllm/entrypoints/openai/serving_utils.h" #include "vllm/entrypoints/openai/tool_parsers/structural_tags.h" #include "vllm/tokenizer/tokenizer.h" @@ -24,6 +29,10 @@ namespace { // self.response_role ("assistant") when add_generation_prompt (the T0 default). constexpr const char* kAssistantRole = "assistant"; +void ChatDbg(const std::string& id, const std::string& msg) { + LogRequestStage(id, msg); +} + // Whether tool_choice selects a single named function (finish_reason stays the // model's own — "stop" — for named calls; chat_completion/serving.py:688,935). bool IsNamedToolChoice(const ChatCompletionRequest& request) { @@ -417,6 +426,18 @@ class ChatSseStream final : public SseStream { previous_num_tokens_ += static_cast(output.token_ids.size()); const std::string current_text = previous_text_ + delta_text; const bool finished = output.finish_reason.has_value() || response.finished; + if (GetRequestLogConfig().debug_stages) { + const auto now = std::chrono::steady_clock::now(); + if (previous_num_tokens_ <= 1 || finished || + std::chrono::duration_cast(now - last_dbg_) + .count() >= 1000) { + ChatDbg(response_id_, + "stage=async_sse prompt_tok=" + std::to_string(prompt_tokens_) + + " gen_tok=" + std::to_string(previous_num_tokens_) + + " finished=" + std::string(finished ? "1" : "0")); + last_dbg_ = now; + } + } std::optional delta = engine_parser_ != nullptr ? ShapeChatDeltaEngine( @@ -494,12 +515,13 @@ class ChatSseStream final : public SseStream { bool role_pending_ = true; bool usage_pending_ = false; bool done_pending_ = false; - bool engine_finished_ = false; bool complete_ = false; + bool engine_finished_ = false; bool aborted_ = false; bool tools_streamed_ = false; int previous_num_tokens_ = 0; std::string previous_text_; + std::chrono::steady_clock::time_point last_dbg_{std::chrono::steady_clock::now()}; }; } // namespace @@ -578,6 +600,56 @@ ChatCompletionResult OpenAIServingChat::create_chat_completion( const std::string prompt = prompt_fn_(request.messages, /*add_generation_prompt=*/true, tools); + const int max_tok_log = + request.max_completion_tokens.has_value() + ? *request.max_completion_tokens + : request.max_tokens.value_or(-1); + std::string roles_summary; + for (size_t i = 0; i < request.messages.size(); ++i) { + if (i) roles_summary += ","; + roles_summary += request.messages[i].role; + size_t clen = request.messages[i].content.has_value() + ? request.messages[i].content->size() + : 0; + roles_summary += "(" + std::to_string(clen) + ")"; + } + LogRequestReceived(request_id, "/v1/chat/completions", model_name, request.stream, + max_tok_log, static_cast(request.messages.size()), + static_cast(tools.size()), prompt.size(), prompt, + roles_summary); + ChatDbg(request_id, "stage=templated prompt_chars=" + std::to_string(prompt.size()) + + " stream=" + std::string(request.stream ? "1" : "0")); + const auto req_t0 = std::chrono::steady_clock::now(); + + // Lab guardrails: Hermes accidentally sending full SOUL (~140k chars) + max_tokens=65536 + // wedges single-batch async prefill for many minutes with no client tokens. + // Override with VT_SERVER_MAX_PROMPT_CHARS / VT_SERVER_MAX_NEW_TOKENS (0 = disable). + static const size_t kMaxPromptChars = [] { + const char* e = std::getenv("VT_SERVER_MAX_PROMPT_CHARS"); + if (e && e[0]) return static_cast(std::strtoull(e, nullptr, 10)); + // Default raised for Hermes full SOUL+tools (~140k). Set lower for safety. + return static_cast(200000); + }(); + static const int kMaxNewTokensCap = [] { + const char* e = std::getenv("VT_SERVER_MAX_NEW_TOKENS"); + if (e && e[0]) return std::atoi(e); + return 4096; // 0 disables + }(); + if (kMaxPromptChars > 0 && prompt.size() > kMaxPromptChars) { + std::ostringstream err; + err << "prompt too large for this server (" << prompt.size() + << " chars > VT_SERVER_MAX_PROMPT_CHARS=" << kMaxPromptChars + << "). Hermes is likely injecting a full system SOUL; shrink the system " + "prompt / tools payload. Set VT_SERVER_MAX_PROMPT_CHARS=0 to disable."; + LogRequestError(request_id, "/v1/chat/completions", err.str()); + throw std::runtime_error(err.str()); + } + if (prompt.size() > 32000) { + ChatDbg(request_id, + "note=large_prompt prefix_caching=ON — first request pays full prefill; " + "identical system+tools prefix on later turns should hit APC"); + } + // ── Multimodal (MM-SERVE-ENGINE) ───────────────────────────────────────── // When the mm seam is set AND a message carries a mm content part, decode + // route the media through the mm processor and carry the placeholder-EXPANDED @@ -676,6 +748,14 @@ ChatCompletionResult OpenAIServingChat::create_chat_completion( const bool named_tool_choice = IsNamedToolChoice(request); SamplingParams sampling_params = request.to_sampling_params(); + if (kMaxNewTokensCap > 0) { + const int before = sampling_params.max_tokens.value_or(kMaxNewTokensCap); + if (before > kMaxNewTokensCap) { + ChatDbg(request_id, "clamp max_tokens " + std::to_string(before) + " -> " + + std::to_string(kMaxNewTokensCap)); + sampling_params.max_tokens = kMaxNewTokensCap; + } + } // tool_choice -> a structural-tag constraint (structured_outputs.structural_tag) // before add_request, built for the ACTIVE tool parser family @@ -692,10 +772,13 @@ ChatCompletionResult OpenAIServingChat::create_chat_completion( const std::string engine_request_id = request_id; if (request.stream) { + ChatDbg(request_id, "stage=stream_begin engine=" + + std::string(async_engine_ != nullptr ? "async" : "sync")); if (async_engine_ != nullptr) { v1::AsyncRequest async_request = async_engine_->add_request( engine_request_id, prompt, std::move(sampling_params), request.priority); + ChatDbg(request_id, "stage=async_queued"); ChatCompletionResult result; result.streaming = true; try { @@ -725,12 +808,28 @@ ChatCompletionResult OpenAIServingChat::create_chat_completion( if (sync_engine_ == nullptr) { throw std::runtime_error("chat handler has no engine"); } + ChatDbg(request_id, "stage=sync_add_request (prefill may take a while on long prompts)"); sync_engine_->add_request(engine_request_id, prompt, std::move(sampling_params), request.priority); + ChatDbg(request_id, "stage=sync_step_loop"); + int step_i = 0; + auto last_prog = std::chrono::steady_clock::now(); while (sync_engine_->has_unfinished_requests()) { for (const RequestOutput& res : sync_engine_->step()) { if (res.request_id != engine_request_id) continue; num_prompt_tokens = static_cast(res.prompt_token_ids.size()); + ++step_i; + const auto now = std::chrono::steady_clock::now(); + if (step_i == 1 || previous_num_tokens == 0 || + std::chrono::duration_cast(now - last_prog) + .count() >= 1000) { + ChatDbg(request_id, + "stage=sync_step n=" + std::to_string(step_i) + + " prompt_tok=" + std::to_string(num_prompt_tokens) + + " gen_tok=" + std::to_string(previous_num_tokens) + + " finished=" + std::string(res.finished ? "1" : "0")); + last_prog = now; + } if (!role_emitted) { role_emitted = true; ChatCompletionResponseStreamChoice role_choice; @@ -829,6 +928,16 @@ ChatCompletionResult OpenAIServingChat::create_chat_completion( "data: " + nlohmann::json(usage_chunk).dump() + "\n\n"); } result.sse_chunks.push_back("data: [DONE]\n\n"); + ChatDbg(request_id, "stage=stream_done prompt_tok=" + + std::to_string(num_prompt_tokens) + + " gen_tok=" + std::to_string(previous_num_tokens)); + { + const double elapsed = + std::chrono::duration(std::chrono::steady_clock::now() - req_t0) + .count(); + LogRequestFinished(request_id, num_prompt_tokens, previous_num_tokens, "stream", + elapsed, previous_text); + } return result; } @@ -836,6 +945,7 @@ ChatCompletionResult OpenAIServingChat::create_chat_completion( // With mm inputs, drive the engine mm overload (placeholder-expanded prompt + // mm_features); otherwise the text-only string overload byte-identically. The // mm forward on the GPU worker consumes the mm_features (MM-SERVE-E2E). + ChatDbg(request_id, "stage=generate_blocking begin (prefill+decode; long prompts stall here)"); const RequestOutput final_res = mm_inputs.has_value() ? (async_engine_ != nullptr @@ -845,11 +955,15 @@ ChatCompletionResult OpenAIServingChat::create_chat_completion( : sync_engine_->generate(std::move(*mm_inputs), std::move(sampling_params), engine_request_id, request.priority)) - : (async_engine_ != nullptr - ? async_engine_->generate(prompt, std::move(sampling_params), - engine_request_id, request.priority) - : sync_engine_->generate(prompt, std::move(sampling_params), - engine_request_id, request.priority)); + : (async_engine_ != nullptr + ? async_engine_->generate(prompt, std::move(sampling_params), + engine_request_id, request.priority) + : sync_engine_->generate(prompt, std::move(sampling_params), + engine_request_id, request.priority)); + ChatDbg(request_id, "stage=generate_blocking done outputs=" + + std::to_string(final_res.outputs.size()) + + " prompt_tok=" + + std::to_string(final_res.prompt_token_ids.size())); ChatCompletionResponse response; response.id = request_id; @@ -900,6 +1014,21 @@ ChatCompletionResult OpenAIServingChat::create_chat_completion( response.usage.completion_tokens = num_generated_tokens; response.usage.total_tokens = num_prompt_tokens + num_generated_tokens; + { + std::string finish = response.choices.empty() + ? "?" + : response.choices[0].finish_reason.value_or("?"); + std::string out_text; + if (!response.choices.empty() && response.choices[0].message.content.has_value()) { + out_text = *response.choices[0].message.content; + } + const double elapsed = + std::chrono::duration(std::chrono::steady_clock::now() - req_t0) + .count(); + LogRequestFinished(request_id, num_prompt_tokens, num_generated_tokens, finish, + elapsed, out_text); + } + ChatCompletionResult result; result.streaming = false; result.response = std::move(response); diff --git a/src/vllm/model_executor/model_loader/nvfp4_dequant.cpp b/src/vllm/model_executor/model_loader/nvfp4_dequant.cpp index b9780ba94..965aa22b0 100644 --- a/src/vllm/model_executor/model_loader/nvfp4_dequant.cpp +++ b/src/vllm/model_executor/model_loader/nvfp4_dequant.cpp @@ -1,12 +1,23 @@ // Ported from: vllm/model_executor/layers/quantization/modelopt.py (NVFP4 W4A16 dequant) @ e24d1b24 #include "vllm/model_executor/model_loader/nvfp4_dequant.h" +#include +#include #include #include +#include +#include #include "vt/dtype.h" namespace vllm { +namespace { +// >0 while expert-level parallel prefetch holds workers (avoid nested storms). +std::atomic g_fp8_dequant_outer_parallel{0}; +} // namespace + +void Fp8DequantBeginOuterParallel() { g_fp8_dequant_outer_parallel.fetch_add(1); } +void Fp8DequantEndOuterParallel() { g_fp8_dequant_outer_parallel.fetch_sub(1); } float F8E4M3ToF32(uint8_t byte) { // IEEE fp8-e4m3fn: 1 sign | 4 exp | 3 mantissa, bias 7, finite (no inf), @@ -90,13 +101,37 @@ void DequantFp8ChannelToBf16(const uint8_t* weight_f8, const uint16_t* scale_bf1 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); + + auto row_work = [&](int64_t n0, int64_t n1) { + for (int64_t n = n0; n < n1; ++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); + } + }; + + // Parallelize over output rows when large enough (MoE expert I/H dims). + // Skip when already under expert-level parallel prefetch. + const int hw = static_cast(std::thread::hardware_concurrency()); + const int nt = (N >= 64 && hw > 1 && g_fp8_dequant_outer_parallel.load() == 0) + ? std::min(hw, 8) + : 1; + if (nt == 1) { + row_work(0, N); + return; + } + std::vector pool; + pool.reserve(static_cast(nt)); + const int64_t chunk = (N + nt - 1) / nt; + for (int t = 0; t < nt; ++t) { + const int64_t n0 = static_cast(t) * chunk; + const int64_t n1 = std::min(N, n0 + chunk); + if (n0 >= n1) break; + pool.emplace_back(row_work, n0, n1); } + for (auto& th : pool) th.join(); } } // namespace vllm diff --git a/src/vllm/model_executor/models/gemma4.cpp b/src/vllm/model_executor/models/gemma4.cpp index cd9d2c44c..a85545b4e 100644 --- a/src/vllm/model_executor/models/gemma4.cpp +++ b/src/vllm/model_executor/models/gemma4.cpp @@ -31,8 +31,11 @@ // residual, not a correctness gap). #include "vllm/model_executor/models/gemma4.h" +#include +#include #include #include +#include #include #include #include @@ -49,6 +52,7 @@ #include "vt/backend.h" #include "vt/dtype.h" #include "vt/ops.h" +#include "vt/fused_ops.h" namespace vllm { namespace { @@ -210,18 +214,37 @@ DBuf Gemma4AttnBlock(Dev d, const Gemma4LayerWeights& w, const Gemma4Layout& g, "gemma4: KV cache head dims mismatch this layer (heterogeneous KV — " "runner must allocate per-layer head_dim; see gemma4.h G1 note)"); - // Merged QKVParallelLinear: D1 folds the shared-input q/k/v GEMMs to ONE - // MatmulBT over the merged [qdim+2kdim,H] owner + a contiguous QkvSplit - // (MergedQkvEnabled(), VT_QWEN3_QKV_MERGE default ON; =0 = byte-identical - // 3-shard). The QKV GEMM is uniform across all Gemma-4 layers — the - // heterogeneous sliding/shared-KV/norm handling downstream is unaffected. - DBuf q(d, adt, {T, qdim}); - DBuf k(d, adt, {T, kdim}); - DBuf v(d, adt, {T, kdim}); + // TLS temps across layers (decode T=1 thrash). + struct AttnTls { + int dev = -1; + int64_t T = 0, qdim = 0, kdim = 0, Hq = 0, Dh = 0; + std::optional q, k, v, qkv, attn; + }; + static thread_local AttnTls tls; + if (tls.dev != d.q.device.index || tls.T != T || tls.qdim != qdim || tls.kdim != kdim || + tls.Hq != Hq || tls.Dh != Dh) { + tls.q.emplace(d, adt, std::vector{T, qdim}); + tls.k.emplace(d, adt, std::vector{T, kdim}); + tls.v.emplace(d, adt, std::vector{T, kdim}); + tls.qkv.emplace(d, adt, std::vector{T, qdim + 2 * kdim}); + tls.attn.emplace(d, adt, std::vector{T, Hq, Dh}); + tls.dev = d.q.device.index; + tls.T = T; + tls.qdim = qdim; + tls.kdim = kdim; + tls.Hq = Hq; + tls.Dh = Dh; + } + DBuf& q = *tls.q; + DBuf& k = *tls.k; + DBuf& v = *tls.v; + DBuf& qkv = *tls.qkv; + DBuf& attn = *tls.attn; + + // Merged QKVParallelLinear: one MatmulBT + QkvSplit (default). { Tensor wqkv = ResidentWeight(d, w.attn.qkv_proj); if (MergedQkvEnabled()) { - DBuf qkv(d, adt, {T, qdim + 2 * kdim}); vt::MatmulBT(d.q, qkv.t(), dhn, wqkv); vt::QkvSplit(d.q, q.t(), k.t(), v.t(), qkv.t()); } else { @@ -273,9 +296,10 @@ DBuf Gemma4AttnBlock(Dev d, const Gemma4LayerWeights& w, const Gemma4Layout& g, Tensor v3 = Reshape(v.t(), {T, Hkv, Dh}); Tensor kw = k3; Tensor vw = v3; - DBuf kcast(d, kv.dtype, {T, Hkv, Dh}); - DBuf vcast(d, kv.dtype, {T, Hkv, Dh}); + // Cast buffers only when dtype differs (rare for bf16 KV). if (kv.dtype != adt) { + DBuf kcast(d, kv.dtype, {T, Hkv, Dh}); + DBuf vcast(d, kv.dtype, {T, Hkv, Dh}); if (kv.dtype == DType::kBF16) { vt::CastBf16(d.q, kcast.t(), k3); vt::CastBf16(d.q, vcast.t(), v3); @@ -285,17 +309,19 @@ DBuf Gemma4AttnBlock(Dev d, const Gemma4LayerWeights& w, const Gemma4Layout& g, } kw = kcast.t(); vw = vcast.t(); + Tensor k_cache = KvSlice(kv, d.q.device, 0); + Tensor v_cache = KvSlice(kv, d.q.device, 1); + vt::ReshapeAndCache(d.q, kw, vw, k_cache, v_cache, si.slot_mapping.t()); + } else { + Tensor k_cache = KvSlice(kv, d.q.device, 0); + Tensor v_cache = KvSlice(kv, d.q.device, 1); + vt::ReshapeAndCache(d.q, kw, vw, k_cache, v_cache, si.slot_mapping.t()); } - Tensor k_cache = KvSlice(kv, d.q.device, 0); - Tensor v_cache = KvSlice(kv, d.q.device, 1); - vt::ReshapeAndCache(d.q, kw, vw, k_cache, v_cache, si.slot_mapping.t()); } - // Paged GQA attention: scale = 1.0 (Q/K norms carry the scale). Reads the - // target layer's populated cache for shared layers. + // Paged GQA attention: scale = 1.0 (Q/K norms carry the scale). Tensor k_cache = KvSlice(kv, d.q.device, 0); Tensor v_cache = KvSlice(kv, d.q.device, 1); - DBuf attn(d, adt, {T, Hq, Dh}); vt::PagedAttentionArgs pa{1.0f, meta.causal}; if (g.attn_logit_softcap > 0.0f) pa.logits_soft_cap = g.attn_logit_softcap; pa.query_start_loc_host = meta.query_start_loc.data(); @@ -307,7 +333,7 @@ DBuf Gemma4AttnBlock(Dev d, const Gemma4LayerWeights& w, const Gemma4Layout& g, Tensor o_in = Reshape(attn.t(), {T, Hq * Dh}); Tensor wo = ResidentWeight(d, w.attn.o_proj); - DBuf o(d, DType::kBF16, {T, H}); + DBuf o(d, DType::kBF16, {T, H}); // returned — not TLS vt::MatmulBT(d.q, o.t(), o_in, wo); return o; } @@ -315,11 +341,25 @@ DBuf Gemma4AttnBlock(Dev d, const Gemma4LayerWeights& w, const Gemma4Layout& g, // GeGLU MLP (gemma4.py::Gemma4MLP). DBuf Gemma4MlpBlock(Dev d, const Gemma4MlpWeights& w, int64_t H, int64_t I, const Tensor& dh2, int64_t T) { - // gate_up MatmulBT -> GeluAndMul(tanh) via the SHARED bf16 GeGLU gate-up MLP seam - // (layers::UnquantizedMlpGateUpGeluMethod). Byte-for-byte the inline sequence — - // folds Gemma-4 onto the shared MlpGateUpMethodBase descriptor. (Tier-C1, - // arch-fusion-fold-plan-2026-07-30.) - DBuf act = layers::UnquantizedMlpGateUpGeluMethod(&w.gate_up_proj, I).Apply(d, dh2); + // Dense GeGLU: TLS reuse of large gate_up [T,2I] + act [T,I] across layers. + Tensor wgu = ResidentWeight(d, w.gate_up_proj); + struct MlpTls { + int dev = -1; + int64_t T = 0, I = 0; + std::optional gu, act; + }; + static thread_local MlpTls tls; + if (tls.dev != d.q.device.index || tls.T != T || tls.I != I) { + tls.gu.emplace(d, DType::kBF16, std::vector{T, 2 * I}); + tls.act.emplace(d, DType::kBF16, std::vector{T, I}); + tls.dev = d.q.device.index; + tls.T = T; + tls.I = I; + } + DBuf& gu = *tls.gu; + DBuf& act = *tls.act; + vt::MatmulBT(d.q, gu.t(), dh2, wgu); + vt::GeluAndMul(d.q, act.t(), gu.t()); Tensor wd = ResidentWeight(d, w.down_proj); DBuf down(d, DType::kBF16, {T, H}); vt::MatmulBT(d.q, down.t(), act.t(), wd); @@ -443,6 +483,24 @@ DBuf ForwardBody(Dev d, const std::vector& token_ids, vt::MulScalar(d.q, ple_input.t(), ple_input.t(), g.input_scale); } + // Layout [L,T,ple] so each layer's slice is one contiguous D2D (not T row gathers). + DBuf ple_by_layer(d, DType::kBF16, + ple > 0 ? std::vector{L, T, ple} : std::vector{1, 1, 1}); + if (ple > 0) { + const size_t row = static_cast(ple) * sizeof(uint16_t); + const auto* src = static_cast(ple_input.ptr()); + auto* dst = static_cast(ple_by_layer.ptr()); + for (int64_t t = 0; t < T; ++t) { + for (int64_t li = 0; li < L; ++li) { + const size_t s_off = + (static_cast(t) * static_cast(L) + static_cast(li)) * row; + const size_t d_off = + (static_cast(li) * static_cast(T) + static_cast(t)) * row; + CopyRow(d, dst + d_off, src + s_off, row); + } + } + } + StepInputs si = BuildStepInputs(d, positions, attn_meta, config); // hidden state stream (each layer fully materializes h; no separate residual). @@ -452,6 +510,53 @@ DBuf ForwardBody(Dev d, const std::vector& token_ids, const size_t ple_row_bytes = static_cast(ple) * sizeof(uint16_t); + // Decode/prefill layer scratch reused across L (pool thrash was real on 30L MoE). + struct LayerTls { + int dev = -1; + int64_t T = 0, H = 0, I = 0, ple = 0; + std::optional dhn, attn_n, h1, dh2, h2, moe_in, n1, n2, sum, n3, mlp_n; + std::optional gate_lin, ple_l, gated, contrib; + }; + static thread_local LayerTls lt; + if (lt.dev != d.q.device.index || lt.T != T || lt.H != H || lt.I != I || lt.ple != ple) { + auto mk = [&](int64_t a, int64_t b) { + return DBuf(d, DType::kBF16, std::vector{a, b}); + }; + lt.dhn.emplace(mk(T, H)); + lt.attn_n.emplace(mk(T, H)); + lt.h1.emplace(mk(T, H)); + lt.dh2.emplace(mk(T, H)); + lt.h2.emplace(mk(T, H)); + lt.moe_in.emplace(mk(T, H)); + lt.n1.emplace(mk(T, H)); + lt.n2.emplace(mk(T, H)); + lt.sum.emplace(mk(T, H)); + lt.n3.emplace(mk(T, H)); + lt.mlp_n.emplace(mk(T, H)); + lt.contrib.emplace(mk(T, H)); + if (ple > 0) { + lt.gate_lin.emplace(mk(T, ple)); + lt.ple_l.emplace(mk(T, ple)); + lt.gated.emplace(mk(T, ple)); + } else { + lt.gate_lin.reset(); + lt.ple_l.reset(); + lt.gated.reset(); + } + lt.dev = d.q.device.index; + lt.T = T; + lt.H = H; + lt.I = I; + lt.ple = ple; + } + DBuf& dhn = *lt.dhn; + DBuf& h1 = *lt.h1; + DBuf& dh2 = *lt.dh2; + DBuf& h2 = *lt.h2; + DBuf& moe_in = *lt.moe_in; + const size_t th_bytes = static_cast(T) * static_cast(H) * sizeof(uint16_t); + DBuf& contrib = *lt.contrib; + for (int64_t l = 0; l < L; ++l) { const Gemma4LayerWeights& w = weights.layers[static_cast(l)]; const int64_t Dh = w.head_dim; @@ -460,101 +565,106 @@ DBuf ForwardBody(Dev d, const std::vector& token_ids, std::optional window; if (!full) window = g.sliding_window; - // Which cache to attend: own for non-shared, target's for YOCO-shared. const int64_t kv_idx = w.is_kv_shared ? w.kv_target_layer : l; VT_CHECK(kv_idx >= 0 && kv_idx < L, "gemma4: bad kv target layer"); const PagedKvCache& kv = attn_kv[static_cast(kv_idx)]; - // r = hidden; dhn = input_layernorm(hidden) [standalone plain] Tensor w_in = ResidentWeight(d, w.input_layernorm, {H}); - DBuf dhn(d, DType::kBF16, {T, H}); vt::RmsNorm(d.q, dhn.t(), hidden.t(), w_in, plain); + static const bool layer_prof = [] { + const char* e = std::getenv("VT_GEMMA4_PROFILE"); + return e && e[0] == '1'; + }(); + using clock = std::chrono::steady_clock; + const auto t0 = layer_prof ? clock::now() : clock::time_point{}; + DBuf attn = Gemma4AttnBlock(d, w, g, dhn.t(), si, attn_meta, kv, T, Dh, full, ones_dh, full ? &prop_cache.t() : nullptr, window, g.rope_theta_sliding); - // attn_n = post_attention_layernorm(attn) [standalone]; hidden = attn_n + r Tensor w_pa = ResidentWeight(d, w.post_attention_layernorm, {H}); - DBuf attn_n(d, DType::kBF16, {T, H}); - vt::RmsNorm(d.q, attn_n.t(), attn.t(), w_pa, plain); - DBuf h1(d, DType::kBF16, {T, H}); - vt::Add(d.q, h1.t(), attn_n.t(), hidden.t()); + // h1 = rmsnorm(attn) + hidden (one kernel) + vt::RmsNormPlusAdd(d.q, h1.t(), attn.t(), w_pa, hidden.t(), plain); + + if (layer_prof) d.b.Synchronize(d.q); + const auto t1 = layer_prof ? clock::now() : clock::time_point{}; - // 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); - DBuf h2(d, DType::kBF16, {T, H}); + if (layer_prof) d.b.Synchronize(d.q); + const auto t2 = layer_prof ? clock::now() : clock::time_point{}; + 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()); + // h2 = rmsnorm(rmsnorm(mlp)+rmsnorm(moe), w_pff) + h1 — one fused kernel + vt::DualRmsNormPlusRes(d.q, h2.t(), mlp.t(), w_p1, moe_out.tensor, w_p2, w_pff, + h1.t(), plain); } 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()); + vt::RmsNormPlusAdd(d.q, h2.t(), mlp.t(), w_pff, h1.t(), plain); + } + + if (layer_prof) { + d.b.Synchronize(d.q); + const auto t3 = clock::now(); + static std::atomic n{0}, us_attn{0}, us_mlp{0}, us_moe{0}; + auto us = [](auto a, auto b) { + return std::chrono::duration_cast(b - a).count(); + }; + us_attn.fetch_add(static_cast(us(t0, t1)), std::memory_order_relaxed); + us_mlp.fetch_add(static_cast(us(t1, t2)), std::memory_order_relaxed); + us_moe.fetch_add(static_cast(us(t2, t3)), std::memory_order_relaxed); + const uint64_t c = n.fetch_add(1, std::memory_order_relaxed) + 1; + if (c == 32 || c % 128 == 0) { + std::fprintf(stderr, + "gemma4 layer profile: calls=%llu attn_us=%.1f mlp_us=%.1f moe_us=%.1f " + "(attn%%=%.0f mlp%%=%.0f moe%%=%.0f)\n", + static_cast(c), + static_cast(us_attn.load()) / c, + static_cast(us_mlp.load()) / c, + static_cast(us_moe.load()) / c, + 100.0 * us_attn.load() / (us_attn.load() + us_mlp.load() + us_moe.load() + 1), + 100.0 * us_mlp.load() / (us_attn.load() + us_mlp.load() + us_moe.load() + 1), + 100.0 * us_moe.load() / (us_attn.load() + us_mlp.load() + us_moe.load() + 1)); + } } - // --- 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}); + DBuf& gate_lin = *lt.gate_lin; + DBuf& ple_l = *lt.ple_l; + DBuf& gated = *lt.gated; vt::MatmulBT(d.q, gate_lin.t(), h2.t(), wg); - DBuf gate_in(d, DType::kBF16, {T, 2 * ple}); // [gate_lin | ple_l] - // Assemble [T, 2*ple]: row t = [gate_lin[t] | ple_input[t, l, :]]. - auto* gi = static_cast(gate_in.ptr()); - const auto* gl = static_cast(gate_lin.ptr()); - const auto* pin = static_cast(ple_input.ptr()); - const size_t two = static_cast(2 * ple) * sizeof(uint16_t); - for (int64_t t = 0; t < T; ++t) { - CopyRow(d, gi + static_cast(t) * two, - gl + static_cast(t) * ple_row_bytes, ple_row_bytes); - const size_t src_off = - (static_cast(t) * static_cast(L) + - static_cast(l)) * - ple_row_bytes; - CopyRow(d, gi + static_cast(t) * two + ple_row_bytes, - pin + src_off, ple_row_bytes); - } - DBuf gated(d, DType::kBF16, {T, ple}); - vt::GeluAndMul(d.q, gated.t(), gate_in.t()); // gelu_tanh(gate_lin)*ple_l + // Contiguous [T,ple] slice for this layer from [L,T,ple] layout. + const size_t layer_bytes = static_cast(T) * ple_row_bytes; + const char* src = static_cast(ple_by_layer.ptr()) + + static_cast(l) * layer_bytes; + d.b.Copy(d.q, ple_l.ptr(), src, layer_bytes); + vt::GeluMulSeparate(d.q, gated.ptr(), gate_lin.ptr(), ple_l.ptr(), T * ple, + DType::kBF16); Tensor wp = ResidentWeight(d, w.per_layer_projection, {H, ple}); - DBuf contrib(d, DType::kBF16, {T, H}); vt::MatmulBT(d.q, contrib.t(), gated.t(), wp); Tensor w_pln = ResidentWeight(d, w.post_per_layer_input_norm, {H}); vt::RmsNorm(d.q, contrib.t(), contrib.t(), w_pln, plain); vt::Add(d.q, h2.t(), h2.t(), contrib.t()); } - // Per-layer learned scalar (gemma4.py:707,765). 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); + d.b.Copy(d.q, hidden.ptr(), h2.ptr(), th_bytes); } // Final norm (plain RMSNorm, standalone — residual is None in vLLM). diff --git a/src/vllm/model_executor/models/gemma4_moe.cpp b/src/vllm/model_executor/models/gemma4_moe.cpp index 7aa34ea8d..42f6133d2 100644 --- a/src/vllm/model_executor/models/gemma4_moe.cpp +++ b/src/vllm/model_executor/models/gemma4_moe.cpp @@ -1,11 +1,16 @@ // 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 #include #include +#include +#include #include #include "vllm/model_executor/model_loader/nvfp4_dequant.h" @@ -14,6 +19,7 @@ #include "vt/backend.h" #include "vt/dtype.h" #include "vt/ops.h" +#include "vt/fused_ops.h" namespace vllm { namespace { @@ -24,54 +30,131 @@ using dense_attn::ResidentWeight; using vt::DType; using vt::Tensor; +// Scratch reused across top-k experts within a token (and host H2D weight slots). +struct ExpertScratch { + DBuf gu; // [T, 2I] fused gate|up activations + DBuf act; // [T, I] + DBuf gu_w; // host-path [2I, H] weight upload + DBuf down_w; + ExpertScratch(Dev d, int64_t T, int64_t I, int64_t H) + : gu(d, DType::kBF16, {T, 2 * I}), + act(d, DType::kBF16, {T, I}), + gu_w(d, DType::kBF16, {2 * I, H}), + down_w(d, DType::kBF16, {H, I}) {} +}; + 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 uint16_t* down_e, int64_t I, int64_t H, ExpertScratch& s) { + 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); + d.b.Copy(d.q, s.gu_w.ptr(), gate_up_e, gu_b); + d.b.Copy(d.q, s.down_w.ptr(), down_e, dn_b); + // One GEMM: x @ W_gu^T -> [T, 2I], then GeluAndMul (interleaved gate|up). + vt::MatmulBT(d.q, s.gu.t(), x, s.gu_w.t()); + vt::GeluAndMul(d.q, s.act.t(), s.gu.t()); + vt::MatmulBT(d.q, out.t(), s.act.t(), s.down_w.t()); +} + +void ExpertGeGLUDeviceAccum(Dev d, DBuf& out, const Tensor& x, const uint16_t* gate_up_e, + const uint16_t* down_e, int64_t I, int64_t H, ExpertScratch& s, + float alpha, float beta) { 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); + const vt::Device dev = d.q.device; + // gate_up_e is contiguous [2I, H] — one BT GEMM instead of two. + Tensor gu_w = + Tensor::Contiguous(const_cast(gate_up_e), DType::kBF16, dev, {2 * I, H}); + vt::MatmulBT(d.q, s.gu.t(), x, gu_w); + vt::GeluAndMul(d.q, s.act.t(), s.gu.t()); + vt::MatmulBTAlphaBeta(d.q, out.ptr(), s.act.ptr(), down_e, static_cast(T), + static_cast(H), static_cast(I), alpha, beta, + DType::kBF16); +} + +void ExpertGeGLUFp8Native(Dev d, DBuf& out, const Tensor& x, const void* fp8_gu, + const void* s_gu, const void* fp8_dn, const void* s_dn, int64_t I, + int64_t H, ExpertScratch& s, float alpha, float beta) { + VT_CHECK(x.shape[0] == 1, "fp8 native: T==1 only"); + vt::MatmulBTFp8Channel(d.q, s.gu.ptr(), x.data, fp8_gu, s_gu, /*M=*/1, + static_cast(2 * I), static_cast(H), 1.f, 0.f); + vt::GeluAndMul(d.q, s.act.t(), s.gu.t()); + if (beta == 0.f) { + vt::MatmulBTFp8Channel(d.q, out.ptr(), s.act.ptr(), fp8_dn, s_dn, /*M=*/1, + static_cast(H), static_cast(I), alpha, 0.f); + } else { + DBuf ytmp(d, DType::kBF16, {1, H}); + vt::MatmulBTFp8Channel(d.q, ytmp.ptr(), s.act.ptr(), fp8_dn, s_dn, /*M=*/1, + static_cast(H), static_cast(I), 1.f, 0.f); + vt::MulScalar(d.q, out.t(), out.t(), static_cast(beta)); + DBuf ysc(d, DType::kBF16, {1, H}); + vt::MulScalar(d.q, ysc.t(), ytmp.t(), static_cast(alpha)); + vt::Add(d.q, out.t(), out.t(), ysc.t()); } - 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]; +// Top-k experts: all gate_up GEMMs → one GeluAndMul → all down GEMMs (alpha/beta mix). +// Cuts (top_k-1) Gelu launches vs per-expert ExpertGeGLUDeviceAccum. +// Uses MatmulBTAlphaBetaRocm directly (no vt::MatmulBT dispatch overhead). +bool ExpertGeGLUTopKFusedGelu(Dev d, DBuf& ysum, const Tensor& x, const uint16_t* const* gu_ptrs, + const uint16_t* const* dn_ptrs, const float* wts, int G, int64_t I, + int64_t H) { + if (G <= 0 || x.shape[0] != 1) return false; + struct Tls { + int dev = -1; + int Gcap = 0; + int64_t I = 0, H = 0; + std::optional gu; // [G, 2I] + std::optional act; // [G, I] + }; + static thread_local Tls tls; + if (tls.dev != d.q.device.index || tls.Gcap < G || tls.I != I || tls.H != H) { + tls.gu.emplace(d, DType::kBF16, std::vector{G, 2 * I}); + tls.act.emplace(d, DType::kBF16, std::vector{G, I}); + tls.dev = d.q.device.index; + tls.Gcap = G; + tls.I = I; + tls.H = H; + } 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); + const size_t gu_row = static_cast(2 * I) * 2; + const size_t act_row = static_cast(I) * 2; + const int Ngu = static_cast(2 * I); + const int Nh = static_cast(H); + const int Ki = static_cast(I); + const int Kh = static_cast(H); + + // Phase 1: gate_up GEMMs into packed [G, 2I] + for (int g = 0; g < G; ++g) { + void* gu_out = static_cast(tls.gu->ptr()) + static_cast(g) * gu_row; + vt::MatmulBTAlphaBeta(d.q, gu_out, x.data, gu_ptrs[g], /*M=*/1, Ngu, Kh, 1.f, 0.f, + DType::kBF16); + } + + // Phase 2: single GeluAndMul over all experts + Tensor gu_all = Tensor::Contiguous(static_cast(tls.gu->ptr()), DType::kBF16, dev, + {G, 2 * I}); + Tensor act_all = Tensor::Contiguous(static_cast(tls.act->ptr()), DType::kBF16, dev, + {G, I}); + vt::GeluAndMul(d.q, act_all, gu_all); + + // Phase 3: down GEMMs with alpha/beta accumulate into ysum + for (int g = 0; g < G; ++g) { + const float alpha = wts[g]; + const float beta = (g == 0) ? 0.f : 1.f; + void* act_g = static_cast(tls.act->ptr()) + static_cast(g) * act_row; + vt::MatmulBTAlphaBeta(d.q, ysum.ptr(), act_g, dn_ptrs[g], /*M=*/1, Nh, Ki, alpha, + beta, DType::kBF16); } - 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); + return true; +} + +// Batched top-k path (gather+strided or pointer-batch): currently disabled. +// Lab: gather+strided produced wrong tokens (~23 t/s); pointer-batch ~0.8 t/s. +// Serial / fused-gelu top-k remains the correct path (~34 t/s). +bool ExpertGeGLUDeviceBatched(Dev /*d*/, DBuf& /*ysum*/, const Tensor& /*x*/, + const std::vector& /*gu_ptrs*/, + const std::vector& /*dn_ptrs*/, + const std::vector& /*wts*/, int64_t /*I*/, int64_t /*H*/) { + return false; } } // namespace @@ -94,30 +177,201 @@ void EnsureGemma4Fp8ExpertCached(const Gemma4Fp8ExpertMats& ex, int64_t I, int64 DequantFp8ChannelToBf16(ex.down_w.bytes.data(), reinterpret_cast(ex.down_s.bytes.data()), H, I, ex.cached_dn.data()); + PinGemma4Fp8ExpertHostCache(ex); +} + +// Dequant into caller buffers without retaining a permanent host BF16 cache. +// Used by dual-GPU resident upload (must not pin ~1.5GiB/layer on host). +void DequantGemma4Fp8ExpertToBf16Ephemeral(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 ephemeral dequant null out"); + 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) { + 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)); + return; + } + DequantFp8ChannelToBf16(ex.gate_w.bytes.data(), + reinterpret_cast(ex.gate_s.bytes.data()), I, H, + gate_up_out); + DequantFp8ChannelToBf16(ex.up_w.bytes.data(), + reinterpret_cast(ex.up_s.bytes.data()), I, H, + gate_up_out + I * H); + DequantFp8ChannelToBf16(ex.down_w.bytes.data(), + reinterpret_cast(ex.down_s.bytes.data()), H, I, + down_out); } // Host BF16 cache + device upload once (subsequent tokens use device GEMM path). +// H2D is async on d.q — later GEMMs on the same stream see the data without a +// device-wide Synchronize (was serializing every expert upload). +// VRAM budget: VT_GEMMA4_EXPERT_VRAM_MB (default 12288). LRU evicts oldest. +namespace { +struct DevExpertLru { + struct Slot { + const Gemma4Fp8ExpertMats* ex = nullptr; + void* gu = nullptr; + void* dn = nullptr; + void* fp8_gu = nullptr; + void* fp8_dn = nullptr; + void* s_gu = nullptr; + void* s_dn = nullptr; + size_t bytes = 0; + uint64_t tick = 0; + }; + std::vector slots; + size_t used = 0; + size_t budget = 0; + uint64_t tick = 1; + int dev = -1; + + size_t BudgetBytes() { + if (budget) return budget; + size_t mb = 0; + if (const char* e = std::getenv("VT_GEMMA4_EXPERT_VRAM_MB")) { + const long v = std::strtol(e, nullptr, 10); + if (v >= 0) mb = static_cast(v); + } + budget = mb == 0 ? static_cast(-1) : mb * 1024ull * 1024ull; + return budget; + } + + void EvictOne(Dev d) { + if (slots.empty()) return; + size_t victim = 0; + for (size_t i = 1; i < slots.size(); ++i) + if (slots[i].tick < slots[victim].tick) victim = i; + Slot s = slots[victim]; + if (s.gu) d.b.Free(s.gu); + if (s.dn) d.b.Free(s.dn); + if (s.fp8_gu) d.b.Free(s.fp8_gu); + if (s.fp8_dn) d.b.Free(s.fp8_dn); + if (s.s_gu) d.b.Free(s.s_gu); + if (s.s_dn) d.b.Free(s.s_dn); + if (s.ex) { + s.ex->dev_gu = nullptr; + s.ex->dev_dn = nullptr; + s.ex->dev_fp8_gu = nullptr; + s.ex->dev_fp8_dn = nullptr; + s.ex->dev_s_gu = nullptr; + s.ex->dev_s_dn = nullptr; + } + used = used >= s.bytes ? used - s.bytes : 0; + slots.erase(slots.begin() + static_cast(victim)); + } + + void Note(const Gemma4Fp8ExpertMats* ex, void* gu, void* dn, size_t bytes, Dev d, + void* fp8_gu = nullptr, void* fp8_dn = nullptr, void* s_gu = nullptr, + void* s_dn = nullptr) { + if (dev != d.q.device.index) { + slots.clear(); + used = 0; + dev = d.q.device.index; + } + const size_t bud = BudgetBytes(); + while (used + bytes > bud && !slots.empty()) EvictOne(d); + if (used + bytes > bud) return; + slots.push_back(Slot{ex, gu, dn, fp8_gu, fp8_dn, s_gu, s_dn, bytes, tick++}); + used += bytes; + } + + void Touch(const Gemma4Fp8ExpertMats* ex) { + for (auto& s : slots) { + if (s.ex == ex) { + s.tick = tick++; + return; + } + } + } +}; + +DevExpertLru& ExpertLru() { + static DevExpertLru lru; + return lru; +} +} // namespace + 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; + if (ex.dev_gu != nullptr && ex.dev_dn != nullptr) { + ExpertLru().Touch(&ex); + 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); + const size_t total = gu_b + dn_b; void* gu = nullptr; void* dn = nullptr; try { + auto& lru = ExpertLru(); + while (lru.used + total > lru.BudgetBytes() && !lru.slots.empty()) lru.EvictOne(d); 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; + lru.Note(&ex, gu, dn, total, d); return true; } catch (...) { if (gu) d.b.Free(gu); if (dn) d.b.Free(dn); - return false; // fall back to host H2D path + static std::atomic fails{0}; + const int n = fails.fetch_add(1) + 1; + if (n == 1 || n % 64 == 0) + std::fprintf(stderr, "gemma4 moe: device expert upload fail #%d (falling back to H2D)\n", + n); + return false; + } +} + +// Upload FP8 weights + channel scales (no BF16 dequant). Half weight VRAM vs BF16 path. +bool EnsureGemma4Fp8NativeOnDevice(Dev d, const Gemma4Fp8ExpertMats& ex, int64_t I, int64_t H) { + if (ex.dev_fp8_gu && ex.dev_fp8_dn && ex.dev_s_gu && ex.dev_s_dn) { + ExpertLru().Touch(&ex); + return true; + } + VT_CHECK(ex.gate_w.HasHostBytes() && ex.up_w.HasHostBytes() && ex.down_w.HasHostBytes(), + "fp8 native: missing weights"); + VT_CHECK(ex.gate_s.HasHostBytes() && ex.up_s.HasHostBytes() && ex.down_s.HasHostBytes(), + "fp8 native: missing scales"); + const size_t gu_b = static_cast(2 * I * H); // u8 + const size_t dn_b = static_cast(H * I); // u8 + const size_t sgu_b = static_cast(2 * I) * 2; // bf16 + const size_t sdn_b = static_cast(H) * 2; // bf16 + const size_t total = gu_b + dn_b + sgu_b + sdn_b; + void *fgu = nullptr, *fdn = nullptr, *sgu = nullptr, *sdn = nullptr; + try { + auto& lru = ExpertLru(); + while (lru.used + total > lru.BudgetBytes() && !lru.slots.empty()) lru.EvictOne(d); + fgu = d.b.Alloc(gu_b); + fdn = d.b.Alloc(dn_b); + sgu = d.b.Alloc(sgu_b); + sdn = d.b.Alloc(sdn_b); + // Pack gate|up FP8 rows + d.b.Copy(d.q, fgu, ex.gate_w.bytes.data(), static_cast(I * H)); + d.b.Copy(d.q, static_cast(fgu) + static_cast(I * H), ex.up_w.bytes.data(), + static_cast(I * H)); + d.b.Copy(d.q, fdn, ex.down_w.bytes.data(), dn_b); + d.b.Copy(d.q, sgu, ex.gate_s.bytes.data(), static_cast(I) * 2); + d.b.Copy(d.q, static_cast(sgu) + static_cast(I) * 2, ex.up_s.bytes.data(), + static_cast(I) * 2); + d.b.Copy(d.q, sdn, ex.down_s.bytes.data(), sdn_b); + ex.dev_fp8_gu = fgu; + ex.dev_fp8_dn = fdn; + ex.dev_s_gu = sgu; + ex.dev_s_dn = sdn; + lru.Note(&ex, nullptr, nullptr, total, d, fgu, fdn, sgu, sdn); + return true; + } catch (...) { + if (fgu) d.b.Free(fgu); + if (fdn) d.b.Free(fdn); + if (sgu) d.b.Free(sgu); + if (sdn) d.b.Free(sdn); + return false; } } @@ -151,11 +405,29 @@ Gemma4MoeScratch RunGemma4Moe(vt::Queue& q, const Gemma4MoeLayerWeights& moe, const vt::RmsNormArgs plain{rms_eps, false}; const int compute_dev = q.device.index; + static const bool profile = [] { + const char* e = std::getenv("VT_GEMMA4_PROFILE"); + return e && e[0] == '1'; + }(); + using clock = std::chrono::steady_clock; + const auto t_all0 = profile ? clock::now() : clock::time_point{}; + DBuf rn(d, DType::kBF16, {T, H}); + // Identity RMS weight (ones) — TLS, upload once (was H2D every layer/token). { - 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); + struct OnesTls { + int dev = -1; + int64_t H = 0; + std::optional w; + }; + static thread_local OnesTls ot; + if (ot.dev != compute_dev || ot.H != H || !ot.w) { + std::vector ones(static_cast(H), vt::F32ToBF16(1.f)); + ot.w.emplace(d, DType::kBF16, std::vector{H}, ones.data()); + ot.dev = compute_dev; + ot.H = H; + } + vt::RmsNorm(d.q, rn.t(), router_in, ot.w->t(), plain); } const OwnedTensor& rproj = @@ -165,16 +437,36 @@ Gemma4MoeScratch RunGemma4Moe(vt::Queue& q, const Gemma4MoeLayerWeights& moe, Tensor wp = ResidentWeight(d, rproj); DBuf logits(d, DType::kF32, {T, E}); vt::MatmulBT(d.q, logits.t(), rn.t(), wp); + + // Device router top-k (softmax + greedy). Only D2H [T,K] weights/indices. + DBuf rw(d, DType::kF32, {T, top_k}); + DBuf ri(d, DType::kI32, {T, top_k}); + vt::MoeRouterTopKArgs rargs; + rargs.top_k = top_k; + rargs.renormalize = true; + vt::MoeRouterTopK(d.q, rw.t(), ri.t(), logits.t(), rargs); + + std::vector hw(static_cast(T * top_k)); + std::vector hi(static_cast(T * top_k)); + d.b.Copy(d.q, hw.data(), rw.ptr(), hw.size() * sizeof(float)); + d.b.Copy(d.q, hi.data(), ri.ptr(), hi.size() * sizeof(int32_t)); 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); + + const auto t_router1 = profile ? clock::now() : clock::time_point{}; 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]); } + // Apply per-expert scale to selected weights. + for (int64_t t = 0; t < T; ++t) { + for (int i = 0; i < top_k; ++i) { + const size_t o = static_cast(t * top_k + i); + const int e = hi[o]; + if (e >= 0 && e < static_cast(E)) hw[o] *= hscale[static_cast(e)]; + } + } const auto& ex = moe.experts; const int64_t gu_stride = 2 * I * H; @@ -191,66 +483,313 @@ Gemma4MoeScratch RunGemma4Moe(vt::Queue& q, const Gemma4MoeLayerWeights& moe, DBuf acc(d, DType::kBF16, {T, H}); acc.Zero(d); + // Reuse MoE decode scratch across layers (30 layers × every token was thrashing the pool). + struct MoeTlsScratch { + int dev = -1; + int64_t I = 0, H = 0; + std::unique_ptr esc; + std::optional xin, ysum, y, ysc; + std::optional gu_sc, dn_sc; + bool have_peer = false; + std::vector gu_tmp, dn_tmp; + }; + static thread_local MoeTlsScratch tls; + if (tls.dev != compute_dev || tls.I != I || tls.H != H) { + tls.esc = std::make_unique(d, /*T=*/1, I, H); + tls.xin.emplace(d, DType::kBF16, std::vector{1, H}); + tls.ysum.emplace(d, DType::kBF16, std::vector{1, H}); + tls.y.emplace(d, DType::kBF16, std::vector{1, H}); + tls.ysc.emplace(d, DType::kBF16, std::vector{1, H}); + tls.gu_sc.reset(); + tls.dn_sc.reset(); + tls.have_peer = false; + tls.gu_tmp.clear(); + tls.dn_tmp.clear(); + tls.dev = compute_dev; + tls.I = I; + tls.H = H; + } + ExpertScratch& esc = *tls.esc; + DBuf& xin = *tls.xin; + DBuf& ysum = *tls.ysum; + DBuf& y = *tls.y; + DBuf& ysc = *tls.ysc; + const bool need_peer_sc = + ex.gate_up_dev != nullptr && ex.down_dev != nullptr && !same_dev; + if (need_peer_sc && !tls.have_peer) { + tls.gu_sc.emplace(d, DType::kBF16, std::vector{2 * I, H}); + tls.dn_sc.emplace(d, DType::kBF16, std::vector{H, I}); + tls.have_peer = true; + } + std::optional& gu_sc = tls.gu_sc; + std::optional& dn_sc = tls.dn_sc; + if (ex.is_fp8 && tls.gu_tmp.size() != static_cast(gu_stride)) { + tls.gu_tmp.resize(static_cast(gu_stride)); + tls.dn_tmp.resize(static_cast(dn_stride)); + } + std::vector& gu_tmp = tls.gu_tmp; + std::vector& dn_tmp = tls.dn_tmp; + static const bool host_axpy = [] { + const char* e = std::getenv("VT_GEMMA4_HOST_AXPY"); + return e && e[0] == '1'; + }(); + static const bool batch_experts = [] { + const char* e = std::getenv("VT_GEMMA4_BATCH_EXPERTS"); + return e && e[0] == '1'; + }(); + static const bool fp8_native = [] { + const char* e = std::getenv("VT_GEMMA4_FP8_NATIVE"); + return e && e[0] == '1'; + }(); + static const bool custom_expert = [] { + const char* e = std::getenv("VT_GEMMA4_CUSTOM_EXPERT"); + return e && e[0] == '1'; + }(); + std::vector hsum; + if (host_axpy) hsum.assign(static_cast(H), vt::F32ToBF16(0.f)); + 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 idx(static_cast(top_k)); 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)]; + const size_t o = static_cast(t * top_k + i); + idx[static_cast(i)] = static_cast(hi[o]); + wts[static_cast(i)] = hw[o]; } - 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); + if (host_axpy) { + std::fill(hsum.begin(), hsum.end(), vt::F32ToBF16(0.f)); + } + // device path: first expert MulScalar writes ysum (no Zero needed) + + // Prefetch BF16 caches for this token's top-k experts in parallel (cold only). + if (ex.is_fp8 && !same_dev && ex.gate_up_dev == nullptr) { + bool any_cold = false; + for (int i = 0; i < top_k; ++i) { + const auto& fex = ex.fp8[static_cast(idx[static_cast(i)])]; + if (fex.cached_gu.empty() || fex.cached_dn.empty()) { + any_cold = true; + break; + } + } + if (any_cold) { + Fp8DequantBeginOuterParallel(); + std::vector pref; + pref.reserve(static_cast(top_k)); + for (int i = 0; i < top_k; ++i) { + const int e = idx[static_cast(i)]; + pref.emplace_back([&, e] { + EnsureGemma4Fp8ExpertCached(ex.fp8[static_cast(e)], I, H); + }); + } + for (auto& th : pref) th.join(); + Fp8DequantEndOuterParallel(); + } + } + + // Prefetch: queue device expert H2D for all top-k before any GEMM (same stream). + // Skip when fused resident packs exist (same_dev or peer) — those are the source of truth + // and VRAM is already tight after full resident upload. + if (ex.is_fp8 && !same_dev && ex.gate_up_dev == nullptr) { + for (int i = 0; i < top_k; ++i) { + const int e = idx[static_cast(i)]; + if (e >= 0 && e < static_cast(E)) { + if (fp8_native) + (void)EnsureGemma4Fp8NativeOnDevice(d, ex.fp8[static_cast(e)], I, H); + else + (void)EnsureGemma4Fp8ExpertOnDevice(d, ex.fp8[static_cast(e)], I, H); + } + } + } + + // Fused top-k ExpertGeGLU (VT_GEMMA4_FUSED_EXPERTS=1). + if (!host_axpy && ex.is_fp8 && T == 1) { + std::vector gu_ptrs, dn_ptrs; + gu_ptrs.reserve(static_cast(top_k)); + dn_ptrs.reserve(static_cast(top_k)); + bool all_dev = true; + for (int i = 0; i < top_k; ++i) { + const int e = idx[static_cast(i)]; + const auto& fex = ex.fp8[static_cast(e)]; + if (!fex.dev_gu || !fex.dev_dn) { + all_dev = false; + break; + } + gu_ptrs.push_back(static_cast(fex.dev_gu)); + dn_ptrs.push_back(static_cast(fex.dev_dn)); + } + if (all_dev && + RunGemma4FusedTopkExpertGeGLU(d.q, ysum.ptr(), xin.ptr(), gu_ptrs.data(), + dn_ptrs.data(), wts.data(), top_k, I, H)) { + d.b.Copy( + d.q, + static_cast(acc.ptr()) + static_cast(t) * static_cast(H) * 2, + ysum.ptr(), static_cast(H) * 2); + continue; + } + } + + // Batched path: VT_GEMMA4_BATCH_EXPERTS=1 (default off). + if (batch_experts && ex.is_fp8 && !host_axpy) { + std::vector gu_ptrs; + std::vector dn_ptrs; + gu_ptrs.reserve(static_cast(top_k)); + dn_ptrs.reserve(static_cast(top_k)); + bool all_dev = true; + for (int i = 0; i < top_k; ++i) { + const int e = idx[static_cast(i)]; + const auto& fex = ex.fp8[static_cast(e)]; + if (!EnsureGemma4Fp8ExpertOnDevice(d, fex, I, H)) { + all_dev = false; + break; + } + gu_ptrs.push_back(static_cast(fex.dev_gu)); + dn_ptrs.push_back(static_cast(fex.dev_dn)); + } + if (all_dev && ExpertGeGLUDeviceBatched(d, ysum, xin.t(), gu_ptrs, dn_ptrs, wts, I, H)) { + d.b.Copy( + d.q, + static_cast(acc.ptr()) + static_cast(t) * static_cast(H) * 2, + ysum.ptr(), static_cast(H) * 2); + continue; // next token + } + } + + // Fused-Gelu top-k: gate_up×G → one GeluAndMul → down×G (default BF16 device path). + // Optional custom RDNA4 expert kernels: VT_GEMMA4_CUSTOM_EXPERT=1 + if (!host_axpy && T == 1 && !fp8_native) { + std::vector gu_p, dn_p; + gu_p.reserve(static_cast(top_k)); + dn_p.reserve(static_cast(top_k)); + bool ok = true; + for (int i = 0; i < top_k && ok; ++i) { + const int e = idx[static_cast(i)]; + if (same_dev) { + gu_p.push_back(static_cast(ex.gate_up_dev) + + static_cast(e) * gu_stride); + dn_p.push_back(static_cast(ex.down_dev) + + static_cast(e) * dn_stride); + } else if (ex.is_fp8) { + const auto& fex = ex.fp8[static_cast(e)]; + if (!fex.dev_gu || !fex.dev_dn) { + if (!EnsureGemma4Fp8ExpertOnDevice(d, fex, I, H)) { + ok = false; + break; + } + } + gu_p.push_back(static_cast(fex.dev_gu)); + dn_p.push_back(static_cast(fex.dev_dn)); + } else { + ok = false; + } + } + if (ok && static_cast(gu_p.size()) == top_k) { + bool ran = false; + if (custom_expert) { + std::vector gu_v(static_cast(top_k)); + std::vector dn_v(static_cast(top_k)); + for (int g = 0; g < top_k; ++g) { + gu_v[static_cast(g)] = gu_p[static_cast(g)]; + dn_v[static_cast(g)] = dn_p[static_cast(g)]; + } + ran = vt::ExpertGeGLUBf16TopKM1(d.q, ysum.ptr(), xin.ptr(), gu_v.data(), + dn_v.data(), wts.data(), top_k, + static_cast(I), static_cast(H)); + } + if (!ran) { + ran = ExpertGeGLUTopKFusedGelu(d, ysum, xin.t(), gu_p.data(), dn_p.data(), wts.data(), + top_k, I, H); + } + if (ran) { + d.b.Copy( + d.q, + static_cast(acc.ptr()) + static_cast(t) * static_cast(H) * 2, + ysum.ptr(), static_cast(H) * 2); + continue; + } + } + } + for (int i = 0; i < top_k; ++i) { const int e = idx[static_cast(i)]; - DBuf y(d, DType::kBF16, {1, H}); + const float ww = wts[static_cast(i)]; + const float beta = (i == 0) ? 0.f : 1.f; + bool fused_mix = false; + 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); + ExpertGeGLUDeviceAccum(d, ysum, xin.t(), gu, dn, I, H, esc, ww, beta); + fused_mix = true; + } else if (fp8_native && ex.is_fp8 && T == 1) { + const auto& fex = ex.fp8[static_cast(e)]; + if (EnsureGemma4Fp8NativeOnDevice(d, fex, I, H)) { + ExpertGeGLUFp8Native(d, ysum, xin.t(), fex.dev_fp8_gu, fex.dev_s_gu, fex.dev_fp8_dn, + fex.dev_s_dn, I, H, esc, ww, beta); + fused_mix = true; + } + } else if (need_peer_sc && gu_sc && dn_sc) { + if (PeerCopyGemma4ExpertSlice(ex.dev_id, ex.gate_up_dev, ex.down_dev, e, I, H, + compute_dev, gu_sc->ptr(), dn_sc->ptr())) { + ExpertGeGLUDeviceAccum(d, ysum, xin.t(), static_cast(gu_sc->ptr()), + static_cast(dn_sc->ptr()), I, H, esc, ww, + beta); + fused_mix = true; + } else if (ex.is_fp8) { + const auto& fex = ex.fp8[static_cast(e)]; + DequantGemma4Fp8ExpertToBf16Ephemeral(fex, I, H, gu_tmp.data(), dn_tmp.data()); + ExpertGeGLUHost(d, y, xin.t(), gu_tmp.data(), dn_tmp.data(), I, H, esc); + } else { + VT_CHECK(gu_host && dn_host, "gemma4 moe: peer fail no host"); + ExpertGeGLUHost(d, y, xin.t(), gu_host + static_cast(e) * gu_stride, + dn_host + static_cast(e) * dn_stride, I, H, esc); + } } 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); + ExpertGeGLUDeviceAccum(d, ysum, xin.t(), static_cast(fex.dev_gu), + static_cast(fex.dev_dn), I, H, esc, ww, + beta); + fused_mix = true; } else { EnsureGemma4Fp8ExpertCached(fex, I, H); - ExpertGeGLUHost(d, y, xin.t(), fex.cached_gu.data(), fex.cached_dn.data(), I, H); + ExpertGeGLUHost(d, y, xin.t(), fex.cached_gu.data(), fex.cached_dn.data(), I, H, + esc); } - } else { + } else if (gu_host && dn_host) { ExpertGeGLUHost(d, y, xin.t(), gu_host + static_cast(e) * gu_stride, - dn_host + static_cast(e) * dn_stride, I, H); + dn_host + static_cast(e) * dn_stride, I, H, esc); + } else { + VT_CHECK(false, "gemma4 moe: no expert weights"); + } + + if (fused_mix) continue; // already accumulated into ysum + + if (host_axpy) { + d.b.Synchronize(d.q); + std::vector hy(static_cast(H)); + d.b.Copy(d.q, hy.data(), y.ptr(), hy.size() * 2); + d.b.Synchronize(d.q); + for (int64_t j = 0; j < H; ++j) + hsum[static_cast(j)] = vt::F32ToBF16( + vt::BF16ToF32(hsum[static_cast(j)]) + + ww * vt::BF16ToF32(hy[static_cast(j)])); + } else if (i == 0) { + vt::MulScalar(d.q, ysum.t(), y.t(), static_cast(ww)); + } else { + vt::MulScalar(d.q, ysc.t(), y.t(), static_cast(ww)); + vt::Add(d.q, ysum.t(), ysum.t(), ysc.t()); } - 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); + } + if (host_axpy) { + d.b.Copy(d.q, ysum.ptr(), hsum.data(), hsum.size() * 2); } d.b.Copy(d.q, static_cast(acc.ptr()) + static_cast(t) * static_cast(H) * 2, @@ -262,6 +801,29 @@ Gemma4MoeScratch RunGemma4Moe(vt::Queue& q, const Gemma4MoeLayerWeights& moe, const size_t alloc = acc.alloc_bytes(); void* p = acc.Release(); r.storage = std::shared_ptr(p, [alloc](void* q) { Pool().Put(alloc, q); }); + + if (profile) { + d.b.Synchronize(d.q); + const auto t_all1 = clock::now(); + static std::atomic ncalls{0}; + static std::atomic us_router{0}; + static std::atomic us_total{0}; + const auto ur = std::chrono::duration_cast(t_router1 - t_all0).count(); + const auto ut = std::chrono::duration_cast(t_all1 - t_all0).count(); + us_router.fetch_add(static_cast(ur), std::memory_order_relaxed); + us_total.fetch_add(static_cast(ut), std::memory_order_relaxed); + const uint64_t c = ncalls.fetch_add(1, std::memory_order_relaxed) + 1; + if (c == 1 || c % 64 == 0) { + const uint64_t tr = us_router.load(std::memory_order_relaxed); + const uint64_t tt = us_total.load(std::memory_order_relaxed); + std::fprintf(stderr, + "gemma4 moe profile: calls=%llu router_us/call=%.1f expert+rest_us/call=%.1f " + "total_us/call=%.1f (router%%=%.0f)\n", + static_cast(c), static_cast(tr) / c, + static_cast(tt - tr) / c, static_cast(tt) / c, + tt ? 100.0 * static_cast(tr) / static_cast(tt) : 0.0); + } + } return r; } @@ -290,6 +852,15 @@ size_t UploadGemma4ExpertsResidentForWeights(Gemma4Weights& weights, "was built without -DVLLM_CPP_HIP; resident preload is a no-op.\n"); return 0; } +bool RunGemma4FusedTopkExpertGeGLU(vt::Queue&, void*, const void*, const uint16_t* const*, + const uint16_t* const*, const float*, int, int64_t, int64_t) { + return false; +} +bool PeerCopyGemma4ExpertSlice(int, const void*, const void*, int, int64_t, int64_t, int, void*, + void*) { + return false; +} +void PinGemma4Fp8ExpertHostCache(const Gemma4Fp8ExpertMats&) {} #endif // VLLM_CPP_HIP } // namespace vllm diff --git a/src/vllm/v1/core/sched/scheduler.cpp b/src/vllm/v1/core/sched/scheduler.cpp index dce84ecfe..46f0c7d80 100644 --- a/src/vllm/v1/core/sched/scheduler.cpp +++ b/src/vllm/v1/core/sched/scheduler.cpp @@ -4,9 +4,13 @@ #include #include +#include #include +#include +#include #include #include +#include #include #include #include @@ -24,6 +28,73 @@ namespace vllm::v1 { namespace { +// Prefill progress (chunked prefill). ON if VT_SERVER_PREFILL_PROGRESS=1 or +// VT_SERVER_VERBOSE=1. Rate-limited ~2 Hz per request. +bool PrefillProgressEnabled() { + static const bool on = [] { + const char* p = std::getenv("VT_SERVER_PREFILL_PROGRESS"); + if (p && p[0] == '0') return false; + if (p && p[0] == '1') return true; + const char* v = std::getenv("VT_SERVER_VERBOSE"); + return v && v[0] == '1'; + }(); + return on; +} + +void MaybeLogPrefillProgress(const Request& request) { + if (!PrefillProgressEnabled()) return; + const int prompt = request.num_prompt_tokens > 0 ? request.num_prompt_tokens + : request.NumTokens(); + if (prompt <= 0) return; + const int computed = request.num_computed_tokens; + // Decode phase: computed exceeds prompt once generation tokens append. + if (computed > prompt && !request.is_prefill_chunk) return; + + using clock = std::chrono::steady_clock; + struct State { + clock::time_point last{}; + int last_computed = -1; + bool logged_done = false; + }; + static std::mutex mu; + static std::unordered_map states; + std::lock_guard lock(mu); + State& st = states[request.request_id]; + const auto now = clock::now(); + const bool done = !request.is_prefill_chunk && computed >= prompt; + if (done) { + if (!st.logged_done) { + std::cerr << "INFO prefill id=" << request.request_id + << " computed=" << std::min(computed, prompt) << "/" << prompt + << " (100%) status=done\n"; + std::cerr.flush(); + st.logged_done = true; + } + if (states.size() > 64) { + for (auto it = states.begin(); it != states.end();) { + if (it->second.logged_done) + it = states.erase(it); + else + ++it; + } + } + return; + } + + const auto ms = + std::chrono::duration_cast(now - st.last).count(); + if (st.last_computed >= 0 && ms < 500 && (computed - st.last_computed) < 2048) { + return; + } + st.last = now; + st.last_computed = computed; + const int shown = std::min(computed, prompt); + const double pct = 100.0 * static_cast(shown) / static_cast(prompt); + std::cerr << "INFO prefill id=" << request.request_id << " computed=" << shown + << "/" << prompt << " (" << pct << "%) status=running\n"; + std::cerr.flush(); +} + // Map the config-level policy onto the request-queue policy. // kLPM (ENG-SGLANG-BEHAVIOR-FLAG) rides on the FCFS deque: the queue mechanics // are identical to fcfs (append / popleft / prepend), and the cache-aware @@ -977,6 +1048,7 @@ void Scheduler::update_after_schedule(SchedulerOutput& scheduler_output) { // needs a grammar bitmask this step. scheduler_output.has_structured_output_requests |= request->use_structured_output() && !request->is_prefill_chunk; + MaybeLogPrefillProgress(*request); } // Flush the finished / preempted id sets (assign fresh sets so the already // copied-out scheduler_output is unaffected). diff --git a/src/vllm/v1/engine/llm_engine.cpp b/src/vllm/v1/engine/llm_engine.cpp index 1d3e43a28..6a7b622ce 100644 --- a/src/vllm/v1/engine/llm_engine.cpp +++ b/src/vllm/v1/engine/llm_engine.cpp @@ -227,37 +227,81 @@ void LLMEngine::abort_request(const std::string& request_id) { RequestOutput LLMEngine::generate(const std::string& prompt, SamplingParams params, const std::string& request_id, int priority) { - // The LLM.generate / _run_engine driver for one request (offline_utils.py:591): - // while self.llm_engine.has_unfinished_requests(): - // for output in self.llm_engine.step(): - // if output.finished: outputs.append(output) + // Offline driver for ONE request. CRITICAL: wait only until *this* request_id + // finishes — not has_unfinished_requests() globally. Otherwise a concurrent + // chat/async job (e.g. huge Hermes SOUL) pins every blocking generate forever. add_request(request_id, prompt, std::move(params), priority); RequestOutput result; - while (has_unfinished_requests()) { + int idle_steps = 0; + int steps = 0; + constexpr int kMaxIdleSteps = 100000; // safety; max_tokens should stop sooner + while (true) { std::vector step_outputs = step(); + ++steps; + bool saw_self = false; + bool self_finished = false; for (RequestOutput& out : step_outputs) { + if (out.request_id != request_id) continue; + saw_self = true; if (out.finished) { result = std::move(out); + self_finished = true; } } + if (self_finished) break; + if (!saw_self) { + ++idle_steps; + if (idle_steps >= kMaxIdleSteps) { + abort_request(request_id); + result.request_id = request_id; + result.finished = true; + break; + } + } else { + idle_steps = 0; + } + // If our request vanished without a finished output, stop. + if (!has_unfinished_requests() && !self_finished) { + break; + } } + (void)steps; return result; } RequestOutput LLMEngine::generate(std::vector prompt_token_ids, SamplingParams params, const std::string& request_id, int priority) { - // TokensPrompt single-request driver (mirrors the string generate loop). add_request(request_id, std::move(prompt_token_ids), std::move(params), priority); RequestOutput result; - while (has_unfinished_requests()) { + int idle_steps = 0; + constexpr int kMaxIdleSteps = 100000; + while (true) { std::vector step_outputs = step(); + bool saw_self = false; + bool self_finished = false; for (RequestOutput& out : step_outputs) { + if (out.request_id != request_id) continue; + saw_self = true; if (out.finished) { result = std::move(out); + self_finished = true; + } + } + if (self_finished) break; + if (!saw_self) { + ++idle_steps; + if (idle_steps >= kMaxIdleSteps) { + abort_request(request_id); + result.request_id = request_id; + result.finished = true; + break; } + } else { + idle_steps = 0; } + if (!has_unfinished_requests() && !self_finished) break; } return result; } @@ -265,19 +309,20 @@ RequestOutput LLMEngine::generate(std::vector prompt_token_ids, RequestOutput LLMEngine::embed(std::vector prompt_token_ids, PoolingParams pooling_params, const std::string& request_id, int priority) { - // The single-request pooling driver (LLM.embed / offline.py:65-119 mirror): - // add the pooling request, then loop step() until it finishes. The finished - // RequestOutput carries the pooled vector in pooling_output. add_pooling_request(request_id, std::move(prompt_token_ids), std::move(pooling_params), priority); RequestOutput result; - while (has_unfinished_requests()) { + while (true) { std::vector step_outputs = step(); + bool self_finished = false; for (RequestOutput& out : step_outputs) { - if (out.finished) { + if (out.request_id == request_id && out.finished) { result = std::move(out); + self_finished = true; } } + if (self_finished) break; + if (!has_unfinished_requests()) break; } return result; } @@ -285,18 +330,35 @@ RequestOutput LLMEngine::embed(std::vector prompt_token_ids, RequestOutput LLMEngine::generate(multimodal::MultiModalInputs mm_inputs, SamplingParams params, const std::string& request_id, int priority) { - // Multimodal single-request driver (mirrors the tokens generate loop). The - // step() forward consumes the carried mm_features on the GPU worker - // (MM-SERVE-E2E) — this driver only proves the loop terminates. add_request(request_id, std::move(mm_inputs), std::move(params), priority); RequestOutput result; - while (has_unfinished_requests()) { + int idle_steps = 0; + constexpr int kMaxIdleSteps = 100000; + while (true) { std::vector step_outputs = step(); + bool saw_self = false; + bool self_finished = false; for (RequestOutput& out : step_outputs) { + if (out.request_id != request_id) continue; + saw_self = true; if (out.finished) { result = std::move(out); + self_finished = true; + } + } + if (self_finished) break; + if (!saw_self) { + ++idle_steps; + if (idle_steps >= kMaxIdleSteps) { + abort_request(request_id); + result.request_id = request_id; + result.finished = true; + break; } + } else { + idle_steps = 0; } + if (!has_unfinished_requests() && !self_finished) break; } return result; } diff --git a/src/vt/fused_ops.cpp b/src/vt/fused_ops.cpp new file mode 100644 index 000000000..b77dbc8f2 --- /dev/null +++ b/src/vt/fused_ops.cpp @@ -0,0 +1,134 @@ +#include "vt/fused_ops.h" + +#include + +#include "vt/dtype.h" + +#if defined(VLLM_CPP_HIP) +#include "vt/rocm/rocm_gelu_mul_sep.h" +#include "vt/rocm/rocm_gemma4_expert_geglu.h" +#include "vt/rocm/rocm_matmul_batch.h" +#include "vt/rocm/rocm_rmsnorm_plus_add.h" +#endif + +namespace vt { + +void RmsNormPlusAdd(Queue& q, Tensor& out, const Tensor& x, const Tensor& w, + const Tensor& addend, const RmsNormArgs& args) { +#if defined(VLLM_CPP_HIP) + if (q.device.type == DeviceType::kROCM) { + rocm::RmsNormPlusAddRocm(q, out, x, w, addend, args); + return; + } +#endif + // Composed reference (CPU / non-ROCm): out = rmsnorm(x) + addend via tmp in out. + // Use out as scratch for rmsnorm then add — requires out != addend alias. + RmsNorm(q, out, x, w, args); + Add(q, out, out, addend); +} + +void DualRmsNormPlusRes(Queue& q, Tensor& out, const Tensor& x1, const Tensor& w1, + const Tensor& x2, const Tensor& w2, const Tensor& w3, + const Tensor& residual, const RmsNormArgs& args) { +#if defined(VLLM_CPP_HIP) + if (q.device.type == DeviceType::kROCM) { + rocm::DualRmsNormPlusResRocm(q, out, x1, w1, x2, w2, w3, residual, args); + return; + } +#endif + // Slow but correct host-side compose using existing ops (allocates temps on device). + Tensor n1 = x1; // shape clone without owning — fall back to sequential RmsNorm+Add + // Prefer throw on non-ROCm discrete GPUs without a known-good compose path. + if (q.device.type != DeviceType::kCPU) { + throw std::runtime_error("vt::DualRmsNormPlusRes: non-ROCm GPU path not registered"); + } + (void)n1; + (void)out; + (void)w1; + (void)x2; + (void)w2; + (void)w3; + (void)residual; + (void)args; + throw std::runtime_error("vt::DualRmsNormPlusRes: CPU compose not yet wired"); +} + +void GeluMulSeparate(Queue& q, void* out, const void* gate, const void* up, int64_t n, + DType dtype) { +#if defined(VLLM_CPP_HIP) + if (q.device.type == DeviceType::kROCM) { + rocm::GeluMulSeparateRocm(q, out, gate, up, n, dtype); + return; + } +#endif + (void)q; + (void)out; + (void)gate; + (void)up; + (void)n; + (void)dtype; + throw std::runtime_error("vt::GeluMulSeparate: ROCm-only fast path in this build"); +} + +void MatmulBTAlphaBeta(Queue& q, void* out, const void* a, const void* b, int M, int N, int K, + float alpha, float beta, DType dtype) { +#if defined(VLLM_CPP_HIP) + if (q.device.type == DeviceType::kROCM) { + rocm::MatmulBTAlphaBetaRocm(q, out, a, b, M, N, K, alpha, beta, dtype); + return; + } +#endif + (void)q; + (void)out; + (void)a; + (void)b; + (void)M; + (void)N; + (void)K; + (void)alpha; + (void)beta; + (void)dtype; + throw std::runtime_error("vt::MatmulBTAlphaBeta: ROCm-only in this build"); +} + +void MatmulBTFp8Channel(Queue& q, void* out, const void* a, const void* b_fp8, + const void* scale_bf16, int M, int N, int K, float alpha, float beta) { +#if defined(VLLM_CPP_HIP) + if (q.device.type == DeviceType::kROCM) { + rocm::MatmulBTFp8ChannelRocm(q, out, a, b_fp8, scale_bf16, M, N, K, alpha, beta); + return; + } +#endif + (void)q; + (void)out; + (void)a; + (void)b_fp8; + (void)scale_bf16; + (void)M; + (void)N; + (void)K; + (void)alpha; + (void)beta; + throw std::runtime_error("vt::MatmulBTFp8Channel: ROCm-only in this build"); +} + +bool ExpertGeGLUBf16TopKM1(Queue& q, void* ysum, const void* x, const void* const* w_gu, + const void* const* w_dn, const float* wts, int G, int I, int H) { +#if defined(VLLM_CPP_HIP) + if (q.device.type == DeviceType::kROCM) { + return rocm::ExpertGeGLUBf16TopKM1Rocm(q, ysum, x, w_gu, w_dn, wts, G, I, H); + } +#endif + (void)q; + (void)ysum; + (void)x; + (void)w_gu; + (void)w_dn; + (void)wts; + (void)G; + (void)I; + (void)H; + return false; +} + +} // namespace vt diff --git a/src/vt/rocm/rocm_dense_basic.hip b/src/vt/rocm/rocm_dense_basic.hip index c815d69b4..50ea9889f 100644 --- a/src/vt/rocm/rocm_dense_basic.hip +++ b/src/vt/rocm/rocm_dense_basic.hip @@ -79,6 +79,18 @@ __global__ void GeluMulK(Tout* out, const Tin* x, int64_t n, int64_t d) { } } +// gate[i] and up[i] already separate (no pack) — MoE ExpertGeGLU hot path. +template +__global__ void GeluMulSepK(Tout* out, const Tin* gate, const Tin* up, int64_t n) { + for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < n; + idx += gridDim.x * blockDim.x) { + const float g = Ld(gate, idx); + const float u = Ld(up, idx); + const float inner = 0.7978845608028654f * (g + 0.044715f * g * g * g); + St(out, idx, 0.5f * g * (1.0f + tanhf(inner)) * u); + } +} + 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; @@ -263,6 +275,24 @@ void GeluAndMulKernelRocm(Queue& q, Tensor& out, const Tensor& x) { Check(hipGetLastError(), "gelu_and_mul"); } +// Separate gate/up GeGLU (no pack). n = T * I elements. +void GeluMulSeparateRocm(Queue& q, void* out, const void* gate, const void* up, int64_t n, + DType dtype) { + if (n == 0) return; + hipStream_t st = AsStream(q); + if (dtype == DType::kBF16) + GeluMulSepK<__hip_bfloat16, __hip_bfloat16><<>>( + static_cast<__hip_bfloat16*>(out), static_cast(gate), + static_cast(up), n); + else if (dtype == DType::kF32) + GeluMulSepK<<>>( + static_cast(out), static_cast(gate), static_cast(up), + n); + else + VT_CHECK(false, "rocm gelu_mul_sep dtype"); + Check(hipGetLastError(), "gelu_mul_sep"); +} + 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; diff --git a/src/vt/rocm/rocm_fp8_channel_gemv.hip b/src/vt/rocm/rocm_fp8_channel_gemv.hip new file mode 100644 index 000000000..844deb0d7 --- /dev/null +++ b/src/vt/rocm/rocm_fp8_channel_gemv.hip @@ -0,0 +1,75 @@ +// FP8 E4M3 weight + BF16 per-row channel scale × BF16 activation GEMV (M=1). +// y[n] = alpha * scale[n] * sum_k x[k]*f8(W[n,k]) + beta * y[n] +// Opt-in via VT_GEMMA4_FP8_NATIVE path in gemma4 MoE. +#include +#include + +#include "vt/device.h" +#include "vt/dtype.h" + +namespace vt::rocm { +namespace { + +// IEEE fp8-e4m3fn (matches vllm::F8E4M3ToF32) +__device__ inline float F8E4M3ToF32(uint8_t byte) { + const uint32_t sign = static_cast(byte >> 7) & 0x1U; + const uint32_t exp = static_cast(byte >> 3) & 0xFU; + const uint32_t mant = static_cast(byte) & 0x7U; + const float sm = sign ? -1.0f : 1.0f; + if (exp == 0xFU && mant == 0x7U) return 0.f; // NaN → 0 in matmul + if (exp == 0U) return sm * (static_cast(mant) * (1.0f / 512.0f)); + const float mantissa = 1.0f + static_cast(mant) * (1.0f / 8.0f); + return sm * ldexpf(mantissa, static_cast(exp) - 7); +} + +// x in LDS; each thread owns output rows +__global__ void Fp8ChannelGemvKernel(__hip_bfloat16* __restrict__ y, + const __hip_bfloat16* __restrict__ x, + const uint8_t* __restrict__ W, // [N,K] + const __hip_bfloat16* __restrict__ scale, // [N] + int N, int K, float alpha, float beta) { + extern __shared__ float smem[]; + float* x_cache = smem; + for (int k = static_cast(threadIdx.x); k < K; k += static_cast(blockDim.x)) { + x_cache[k] = __bfloat162float(x[k]); + } + __syncthreads(); + + for (int n = static_cast(blockIdx.x * blockDim.x + threadIdx.x); n < N; + n += static_cast(gridDim.x * blockDim.x)) { + const uint8_t* wrow = W + static_cast(n) * static_cast(K); + float acc = 0.f; + int k = 0; + for (; k + 3 < K; k += 4) { + acc += x_cache[k] * F8E4M3ToF32(wrow[k]); + acc += x_cache[k + 1] * F8E4M3ToF32(wrow[k + 1]); + acc += x_cache[k + 2] * F8E4M3ToF32(wrow[k + 2]); + acc += x_cache[k + 3] * F8E4M3ToF32(wrow[k + 3]); + } + for (; k < K; ++k) acc += x_cache[k] * F8E4M3ToF32(wrow[k]); + const float s = __bfloat162float(scale[n]); + float v = alpha * s * acc; + if (beta != 0.f) v += beta * __bfloat162float(y[n]); + y[n] = __float2bfloat16(v); + } +} + +} // namespace + +// out[1,N] bf16, a[1,K] bf16, b_fp8[N,K], scale_bf16[N] +void MatmulBTFp8ChannelRocm(Queue& q, void* out, const void* a, const void* b_fp8, + const void* scale_bf16, int M, int N, int K, float alpha, + float beta) { + if (M != 1 || N <= 0 || K <= 0) return; + hipStream_t s = static_cast(q.handle); + constexpr int kBlock = 256; + const int grid = (N + kBlock - 1) / kBlock; + const size_t shmem = static_cast(K) * sizeof(float); + if (shmem > 48 * 1024) return; // caller must fall back + Fp8ChannelGemvKernel<<>>( + static_cast<__hip_bfloat16*>(out), static_cast(a), + static_cast(b_fp8), static_cast(scale_bf16), N, K, + alpha, beta); +} + +} // namespace vt::rocm diff --git a/src/vt/rocm/rocm_gemma4_expert_geglu.hip b/src/vt/rocm/rocm_gemma4_expert_geglu.hip new file mode 100644 index 000000000..a4b896a01 --- /dev/null +++ b/src/vt/rocm/rocm_gemma4_expert_geglu.hip @@ -0,0 +1,133 @@ +// Gemma-4 MoE Expert GeGLU decode (T=1), custom RDNA4 HIP kernels. +// VT_GEMMA4_CUSTOM_EXPERT=1 to use from gemma4_moe.cpp. +// y = beta*y + alpha * Down(Gelu(Gate(x)) * Up(x)) +#include +#include + +#include + +#include "vt/device.h" +#include "vt/dtype.h" + +namespace vt::rocm { +namespace { + +constexpr int kBlock = 256; + +__device__ inline float gelu_tanh(float g) { + const float inner = 0.7978845608028654f * (g + 0.044715f * g * g * g); + return 0.5f * g * (1.f + tanhf(inner)); +} + +// Block-level float sum (kBlock == 256). +__device__ inline float BlockSum256(float v) { + __shared__ float sm[256]; + const int tid = static_cast(threadIdx.x); + sm[tid] = v; + __syncthreads(); +#pragma unroll + for (int s = 128; s > 0; s >>= 1) { + if (tid < s) sm[tid] += sm[tid + s]; + __syncthreads(); + } + return sm[0]; +} + +// grid: I — act[i] = gelu(dot(x,Wgu[i])) * dot(x,Wgu[I+i]) +__global__ void ExpertActKernel(float* __restrict__ act, const __hip_bfloat16* __restrict__ x, + const __hip_bfloat16* __restrict__ w_gu, int I, int H) { + const int i = static_cast(blockIdx.x); + if (i >= I) return; + + const __hip_bfloat16* grow = w_gu + static_cast(i) * H; + const __hip_bfloat16* urow = w_gu + static_cast(I + i) * H; + + float gate = 0.f, up = 0.f; + // Vectorize by 2 when H even + const int tid = static_cast(threadIdx.x); + for (int h = tid; h < H; h += kBlock) { + const float xv = __bfloat162float(x[h]); + gate += xv * __bfloat162float(grow[h]); + up += xv * __bfloat162float(urow[h]); + } + gate = BlockSum256(gate); + up = BlockSum256(up); + if (tid == 0) act[i] = gelu_tanh(gate) * up; +} + +// grid: ceil(H / kBlock) — y[h] = beta*y[h] + alpha * dot(act, Wdn[h]) +__global__ void ExpertDownKernel(__hip_bfloat16* __restrict__ y, const float* __restrict__ act, + const __hip_bfloat16* __restrict__ w_dn, int I, int H, float alpha, + float beta) { + const int h = static_cast(blockIdx.x * blockDim.x + threadIdx.x); + if (h >= H) return; + const __hip_bfloat16* drow = w_dn + static_cast(h) * I; + float acc = 0.f; + int i = 0; + for (; i + 3 < I; i += 4) { + acc += act[i] * __bfloat162float(drow[i]); + acc += act[i + 1] * __bfloat162float(drow[i + 1]); + acc += act[i + 2] * __bfloat162float(drow[i + 2]); + acc += act[i + 3] * __bfloat162float(drow[i + 3]); + } + for (; i < I; ++i) acc += act[i] * __bfloat162float(drow[i]); + float v = alpha * acc; + if (beta != 0.f) v += beta * __bfloat162float(y[h]); + y[h] = __float2bfloat16(v); +} + +struct ActBuf { + float* p = nullptr; + int cap = 0; +}; + +ActBuf& TlsAct() { + static thread_local ActBuf b; + return b; +} + +bool EnsureAct(int I) { + auto& b = TlsAct(); + if (b.cap >= I && b.p) return true; + if (b.p) (void)hipFree(b.p); + b.p = nullptr; + b.cap = 0; + if (hipMalloc(&b.p, static_cast(I) * sizeof(float)) != hipSuccess) return false; + b.cap = I; + return true; +} + +} // namespace + +// Returns false if launch/setup failed (caller falls back to hipBLAS path). +bool ExpertGeGLUBf16M1Rocm(Queue& q, void* y, const void* x, const void* w_gu, const void* w_dn, + int I, int H, float alpha, float beta) { + if (!y || !x || !w_gu || !w_dn || I <= 0 || H <= 0) return false; + if (!EnsureAct(I)) return false; + + hipStream_t st = static_cast(q.handle); + auto* act = TlsAct().p; + + ExpertActKernel<<>>( + act, static_cast(x), static_cast(w_gu), I, H); + + const int grid_dn = (H + kBlock - 1) / kBlock; + ExpertDownKernel<<>>( + static_cast<__hip_bfloat16*>(y), act, static_cast(w_dn), I, H, alpha, + beta); + + return hipGetLastError() == hipSuccess; +} + +// Top-k: sequential custom experts with alpha/beta mix into ysum. +bool ExpertGeGLUBf16TopKM1Rocm(Queue& q, void* ysum, const void* x, const void* const* w_gu, + const void* const* w_dn, const float* wts, int G, int I, int H) { + if (G <= 0 || !ysum || !x || !w_gu || !w_dn || !wts) return false; + for (int g = 0; g < G; ++g) { + const float beta = (g == 0) ? 0.f : 1.f; + if (!ExpertGeGLUBf16M1Rocm(q, ysum, x, w_gu[g], w_dn[g], I, H, wts[g], beta)) return false; + } + return true; +} + +} // namespace vt::rocm diff --git a/src/vt/rocm/rocm_gemma4_experts.hip b/src/vt/rocm/rocm_gemma4_experts.hip index 8f274cd69..56460d715 100644 --- a/src/vt/rocm/rocm_gemma4_experts.hip +++ b/src/vt/rocm/rocm_gemma4_experts.hip @@ -1,6 +1,9 @@ // ROCm upload of Gemma-4 MoE expert stacks (BF16 fused or FP8→BF16 dequant). +// Packs GPU0 first (compute device), then GPU1. FP8 streams per-expert to +// avoid ~30G host OOM from permanent BF16 caches / full-layer host buffers. #include +#include #include #include #include @@ -18,6 +21,86 @@ void Check(hipError_t e, const char* w) { throw std::runtime_error(std::string("gemma4 resident: ") + w + ": " + hipGetErrorString(e)); } +bool DeviceHasRoom(int dev, size_t need_bytes) { + Check(hipSetDevice(dev), "hipSetDevice"); + size_t free_b = 0, tot_b = 0; + if (hipMemGetInfo(&free_b, &tot_b) != hipSuccess) return true; + // GPU0 is compute: dense weights + KV + activations need real headroom. + // 12 GiB: prior 10 GiB still OOM'd gelu after 42 GiB dual-GPU resident pack. + const size_t headroom = (dev == 0) ? (12ull << 30) : (512ull << 20); + return free_b >= need_bytes + headroom; +} + +bool AllocLayerDev(int dev, size_t gu_bytes, size_t dn_bytes, void** gu, void** dn) { + if (!DeviceHasRoom(dev, gu_bytes + dn_bytes)) return false; + Check(hipSetDevice(dev), "hipSetDevice"); + *gu = nullptr; + *dn = nullptr; + if (hipMalloc(gu, gu_bytes) != hipSuccess) return false; + if (hipMalloc(dn, dn_bytes) != hipSuccess) { + (void)hipFree(*gu); + *gu = nullptr; + return false; + } + return true; +} + +bool UploadFp8LayerStreaming(Gemma4FusedExperts& ex, int dev) { + const int64_t E = ex.num_experts; + const int64_t I = ex.intermediate; + const int64_t H = ex.hidden; + const size_t gu_one = static_cast(2 * I * H) * sizeof(uint16_t); + const size_t dn_one = static_cast(H * I) * sizeof(uint16_t); + const size_t gu_bytes = gu_one * static_cast(E); + const size_t dn_bytes = dn_one * static_cast(E); + + void* gu = nullptr; + void* dn = nullptr; + if (!AllocLayerDev(dev, gu_bytes, dn_bytes, &gu, &dn)) return false; + + std::vector gu_tmp(static_cast(2 * I * H)); + std::vector dn_tmp(static_cast(H * I)); + for (int64_t e = 0; e < E; ++e) { + DequantGemma4Fp8ExpertToBf16Ephemeral(ex.fp8[static_cast(e)], I, H, gu_tmp.data(), + dn_tmp.data()); + // Drop any accidental permanent cache on this expert (decode may have set it). + ex.fp8[static_cast(e)].cached_gu.clear(); + ex.fp8[static_cast(e)].cached_gu.shrink_to_fit(); + ex.fp8[static_cast(e)].cached_dn.clear(); + ex.fp8[static_cast(e)].cached_dn.shrink_to_fit(); + + Check(hipMemcpy(static_cast(gu) + static_cast(e) * gu_one, gu_tmp.data(), + gu_one, hipMemcpyHostToDevice), + "H2D gu e"); + Check(hipMemcpy(static_cast(dn) + static_cast(e) * dn_one, dn_tmp.data(), + dn_one, hipMemcpyHostToDevice), + "H2D dn e"); + } + ex.gate_up_dev = gu; + ex.down_dev = dn; + ex.dev_id = dev; + return true; +} + +bool UploadBf16Layer(Gemma4FusedExperts& ex, int dev) { + 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; + VT_CHECK(ex.gate_up.bytes.size() == gu_bytes && ex.down.bytes.size() == dn_bytes, + "gemma4 resident: bf16 fused size"); + void* gu = nullptr; + void* dn = nullptr; + if (!AllocLayerDev(dev, gu_bytes, dn_bytes, &gu, &dn)) return false; + Check(hipMemcpy(gu, ex.gate_up.bytes.data(), gu_bytes, hipMemcpyHostToDevice), "H2D gu"); + Check(hipMemcpy(dn, ex.down.bytes.data(), dn_bytes, hipMemcpyHostToDevice), "H2D dn"); + ex.gate_up_dev = gu; + ex.down_dev = dn; + ex.dev_id = dev; + return true; +} + } // namespace size_t UploadGemma4ExpertsResident(std::vector& layers, @@ -28,12 +111,30 @@ size_t UploadGemma4ExpertsResident(std::vector& layers, if (available < 1) return 0; if (num_gpus > available) num_gpus = available; + // Enable peer access so off-device resident experts can D2D into compute GPU. + for (int a = 0; a < num_gpus; ++a) { + for (int b = 0; b < num_gpus; ++b) { + if (a == b) continue; + int can = 0; + if (hipDeviceCanAccessPeer(&can, a, b) == hipSuccess && can) { + Check(hipSetDevice(a), "peer set"); + hipError_t pe = hipDeviceEnablePeerAccess(b, 0); + if (pe != hipSuccess && pe != hipErrorPeerAccessAlreadyEnabled) { + std::fprintf(stderr, "gemma4 moe: peer access %d->%d failed: %s\n", a, b, + hipGetErrorString(pe)); + } + } + } + } + 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)); + int fill_dev = 0; + for (size_t li = 0; li < layers.size(); ++li) { if (ok_layers >= max_layers) break; auto& moe = layers[li]; @@ -42,56 +143,38 @@ size_t UploadGemma4ExpertsResident(std::vector& layers, 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 size_t layer_bytes = + static_cast(E * 2 * I * H) * 2 + static_cast(E * H * I) * 2; + + bool ok = false; + if (ex.is_fp8) + ok = UploadFp8LayerStreaming(ex, fill_dev); + else + ok = UploadBf16Layer(ex, fill_dev); - 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)); + if (!ok && fill_dev + 1 < num_gpus) { + ++fill_dev; + std::fprintf(stderr, "gemma4 moe: GPU%d full after %d layers; continuing on GPU%d\n", + fill_dev - 1, ok_layers, fill_dev); + if (ex.is_fp8) + ok = UploadFp8LayerStreaming(ex, fill_dev); + else + ok = UploadBf16Layer(ex, fill_dev); + } + if (!ok) { + std::fprintf(stderr, "gemma4 moe: resident upload stopped at layer %zu (OOM)\n", li); 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; + total += layer_bytes; ++ok_layers; + std::fprintf(stderr, "gemma4 moe: resident layer %zu -> gpu %d (%.2f GiB cumulative)\n", li, + ex.dev_id, total / (1024.0 * 1024.0 * 1024.0)); + std::fflush(stderr); } Check(hipSetDevice(0), "hipSetDevice0"); - std::fprintf(stderr, "gemma4 moe: resident experts gpus=%d layers=%d %.2f GiB\n", num_gpus, + std::fprintf(stderr, "gemma4 moe: RESIDENT DONE gpus=%d layers=%d %.2f GiB\n", num_gpus, ok_layers, total / (1024.0 * 1024.0 * 1024.0)); + std::fflush(stderr); return total; } @@ -108,4 +191,54 @@ size_t UploadGemma4ExpertsResidentForWeights(Gemma4Weights& weights, int num_gpu return n; } +bool PeerCopyGemma4ExpertSlice(int src_dev, const void* gate_up_base, const void* down_base, + int expert_id, int64_t I, int64_t H, int compute_dev, + void* gate_up_dst, void* down_dst) { + if (!gate_up_base || !down_base || !gate_up_dst || !down_dst) return false; + if (src_dev < 0 || compute_dev < 0) return false; + const size_t gu_one = static_cast(2 * I * H) * sizeof(uint16_t); + const size_t dn_one = static_cast(H * I) * sizeof(uint16_t); + const char* gu_src = + static_cast(gate_up_base) + static_cast(expert_id) * gu_one; + const char* dn_src = + static_cast(down_base) + static_cast(expert_id) * dn_one; + if (src_dev == compute_dev) { + Check(hipSetDevice(compute_dev), "peer same set"); + Check(hipMemcpy(gate_up_dst, gu_src, gu_one, hipMemcpyDeviceToDevice), "D2D gu"); + Check(hipMemcpy(down_dst, dn_src, dn_one, hipMemcpyDeviceToDevice), "D2D dn"); + return true; + } + hipError_t e1 = + hipMemcpyPeer(gate_up_dst, compute_dev, const_cast(gu_src), src_dev, gu_one); + hipError_t e2 = + hipMemcpyPeer(down_dst, compute_dev, const_cast(dn_src), src_dev, dn_one); + if (e1 != hipSuccess || e2 != hipSuccess) { + // Fallback: staged through host (still avoids FP8 dequant) + std::vector gu_tmp(gu_one / 2), dn_tmp(dn_one / 2); + Check(hipSetDevice(src_dev), "peer fb src"); + Check(hipMemcpy(gu_tmp.data(), gu_src, gu_one, hipMemcpyDeviceToHost), "D2H gu"); + Check(hipMemcpy(dn_tmp.data(), dn_src, dn_one, hipMemcpyDeviceToHost), "D2H dn"); + Check(hipSetDevice(compute_dev), "peer fb dst"); + Check(hipMemcpy(gate_up_dst, gu_tmp.data(), gu_one, hipMemcpyHostToDevice), "H2D gu"); + Check(hipMemcpy(down_dst, dn_tmp.data(), dn_one, hipMemcpyHostToDevice), "H2D dn"); + } + Check(hipSetDevice(compute_dev), "peer restore compute"); + return true; +} + +void PinGemma4Fp8ExpertHostCache(const Gemma4Fp8ExpertMats& ex) { + if (ex.host_pinned) return; + if (ex.cached_gu.empty() || ex.cached_dn.empty()) return; + const size_t gu_b = ex.cached_gu.size() * sizeof(uint16_t); + const size_t dn_b = ex.cached_dn.size() * sizeof(uint16_t); + hipError_t e1 = + hipHostRegister(const_cast(ex.cached_gu.data()), gu_b, hipHostRegisterDefault); + hipError_t e2 = + hipHostRegister(const_cast(ex.cached_dn.data()), dn_b, hipHostRegisterDefault); + if ((e1 == hipSuccess || e1 == hipErrorHostMemoryAlreadyRegistered) && + (e2 == hipSuccess || e2 == hipErrorHostMemoryAlreadyRegistered)) { + ex.host_pinned = true; + } +} + } // namespace vllm diff --git a/src/vt/rocm/rocm_gemma4_fused_experts.hip b/src/vt/rocm/rocm_gemma4_fused_experts.hip new file mode 100644 index 000000000..ce824a2f1 --- /dev/null +++ b/src/vt/rocm/rocm_gemma4_fused_experts.hip @@ -0,0 +1,150 @@ +// Fused top-k ExpertGeGLU for Gemma-4 MoE decode (T=1). +// Opt-in: VT_GEMMA4_FUSED_EXPERTS=1 +// Two-phase multi-block: (1) act[i]=gelu(gate)*up (2) ysum += w * act@Wdn +#include +#include + +#include + +#include "vt/device.h" + +namespace vllm { +namespace { + +constexpr int kMaxK = 8; +constexpr int kBlock = 256; + +__device__ inline float gelu_tanh(float x) { + const float c = 0.7978845608028654f; + return 0.5f * x * (1.f + tanhf(c * (x + 0.044715f * x * x * x))); +} + +// grid: (I, G) — one block per (expert, intermediate index) +__global__ void ExpertActKernel(__hip_bfloat16* __restrict__ act, // [G, I] + const __hip_bfloat16* __restrict__ x, // [H] + const __hip_bfloat16* const* __restrict__ gu_ptrs, int I, + int H) { + const int i = static_cast(blockIdx.x); + const int e = static_cast(blockIdx.y); + if (i >= I) return; + const __hip_bfloat16* gu = gu_ptrs[e]; + const __hip_bfloat16* grow = gu + static_cast(i) * H; + const __hip_bfloat16* urow = gu + static_cast(I + i) * H; + + float gate = 0.f, up = 0.f; + for (int h = static_cast(threadIdx.x); h < H; h += kBlock) { + const float xv = __bfloat162float(x[h]); + gate += xv * __bfloat162float(grow[h]); + up += xv * __bfloat162float(urow[h]); + } + __shared__ float sg[256], su[256]; + sg[threadIdx.x] = gate; + su[threadIdx.x] = up; + __syncthreads(); + for (int s = 128; s > 0; s >>= 1) { + if (static_cast(threadIdx.x) < s) { + sg[threadIdx.x] += sg[threadIdx.x + s]; + su[threadIdx.x] += su[threadIdx.x + s]; + } + __syncthreads(); + } + if (threadIdx.x == 0) { + act[static_cast(e) * I + i] = __float2bfloat16(gelu_tanh(sg[0]) * su[0]); + } +} + +// grid: (H,) — accumulate all experts into ysum +__global__ void ExpertDownMixKernel(__hip_bfloat16* __restrict__ ysum, + const __hip_bfloat16* __restrict__ act, // [G,I] + const __hip_bfloat16* const* __restrict__ dn_ptrs, + const float* __restrict__ wts, int G, int I, int H) { + const int h = static_cast(blockIdx.x * blockDim.x + threadIdx.x); + if (h >= H) return; + float acc = 0.f; + for (int e = 0; e < G; ++e) { + const __hip_bfloat16* drow = dn_ptrs[e] + static_cast(h) * I; + const __hip_bfloat16* arow = act + static_cast(e) * I; + float dot = 0.f; + for (int i = 0; i < I; ++i) { + dot += __bfloat162float(arow[i]) * __bfloat162float(drow[i]); + } + acc += wts[e] * dot; + } + ysum[h] = __float2bfloat16(acc); +} + +int FusedEnabled() { + static const int on = [] { + if (const char* e = std::getenv("VT_GEMMA4_FUSED_EXPERTS")) return e[0] == '1' ? 1 : 0; + return 0; + }(); + return on; +} + +struct BufTls { + void** d_gu = nullptr; + void** d_dn = nullptr; + float* d_wts = nullptr; + __hip_bfloat16* d_act = nullptr; + int act_cap = 0; +}; + +BufTls& TLS() { + static thread_local BufTls t; + return t; +} + +} // namespace + +bool RunGemma4FusedTopkExpertGeGLU(vt::Queue& q, void* ysum, const void* x, + const uint16_t* const* gu_ptrs, + const uint16_t* const* dn_ptrs, const float* wts, int G, + int64_t I, int64_t H) { + if (!FusedEnabled() || G <= 0 || G > kMaxK || I <= 0 || H <= 0) return false; + if (!ysum || !x || !gu_ptrs || !dn_ptrs || !wts) return false; + + auto& tls = TLS(); + const size_t act_need = static_cast(G) * static_cast(I); + if (tls.act_cap < static_cast(act_need)) { + if (tls.d_act) (void)hipFree(tls.d_act); + if (hipMalloc(&tls.d_act, act_need * sizeof(__hip_bfloat16)) != hipSuccess) return false; + tls.act_cap = static_cast(act_need); + } + if (!tls.d_gu) { + if (hipMalloc(&tls.d_gu, sizeof(void*) * kMaxK) != hipSuccess) return false; + if (hipMalloc(&tls.d_dn, sizeof(void*) * kMaxK) != hipSuccess) return false; + if (hipMalloc(&tls.d_wts, sizeof(float) * kMaxK) != hipSuccess) return false; + } + + hipStream_t st = static_cast(q.handle); + const void* h_gu[kMaxK]; + const void* h_dn[kMaxK]; + for (int i = 0; i < G; ++i) { + h_gu[i] = gu_ptrs[i]; + h_dn[i] = dn_ptrs[i]; + } + if (hipMemcpyAsync(tls.d_gu, h_gu, sizeof(void*) * static_cast(G), + hipMemcpyHostToDevice, st) != hipSuccess) + return false; + if (hipMemcpyAsync(tls.d_dn, h_dn, sizeof(void*) * static_cast(G), + hipMemcpyHostToDevice, st) != hipSuccess) + return false; + if (hipMemcpyAsync(tls.d_wts, wts, sizeof(float) * static_cast(G), hipMemcpyHostToDevice, + st) != hipSuccess) + return false; + + dim3 grid1(static_cast(I), static_cast(G)); + ExpertActKernel<<>>( + tls.d_act, static_cast(x), + reinterpret_cast(tls.d_gu), static_cast(I), + static_cast(H)); + + const int grid2 = static_cast((H + kBlock - 1) / kBlock); + ExpertDownMixKernel<<>>( + static_cast<__hip_bfloat16*>(ysum), tls.d_act, + reinterpret_cast(tls.d_dn), tls.d_wts, G, static_cast(I), + static_cast(H)); + return true; +} + +} // namespace vllm diff --git a/src/vt/rocm/rocm_matmul_hipblaslt.hip b/src/vt/rocm/rocm_matmul_hipblaslt.hip index 01c8ddf5a..70b21f239 100644 --- a/src/vt/rocm/rocm_matmul_hipblaslt.hip +++ b/src/vt/rocm/rocm_matmul_hipblaslt.hip @@ -4,14 +4,21 @@ // // 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 #include #include #include +#include #include +#include #include "vt/ops.h" @@ -54,19 +61,30 @@ struct BlasCtx { hipblasHandle_t handle = nullptr; }; +// Per-thread handle: hipBLAS is not free-threaded; decode is single-threaded per +// engine worker. Avoid global mutex + SetStream on every GEMM (100s×/token). 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); + struct Tls { + int dev = -1; + hipStream_t stream = nullptr; + hipblasHandle_t handle = nullptr; + }; + static thread_local Tls tls; + if (tls.handle == nullptr || tls.dev != device) { + if (tls.handle) { + // device switch rare — destroy old handle + (void)hipblasDestroy(tls.handle); + tls.handle = nullptr; + } + CheckBlas(hipblasCreate(&tls.handle), "hipblasCreate"); + tls.dev = device; + tls.stream = nullptr; + } + if (tls.stream != stream) { + CheckBlas(hipblasSetStream(tls.handle, stream), "hipblasSetStream"); + tls.stream = stream; } - CheckBlas(hipblasSetStream(it->second.handle, stream), "hipblasSetStream"); - return it->second; + return BlasCtx{tls.handle}; } std::string ComboName(const Tensor& a, const Tensor& b, const Tensor& out) { @@ -83,6 +101,241 @@ hipDataType ToBlasType(DType dt) { } } +// BF16 default: COMPUTE_32F (FAST_16BF returns NOT_SUPPORTED on gfx1201 for our BT shapes). +// Override: VT_ROCM_GEMM_COMPUTE=32f | 16bf | 16f +hipblasComputeType_t GemmCompute(DType dt) { + if (dt != DType::kBF16 && dt != DType::kF16) return HIPBLAS_COMPUTE_32F; + static const hipblasComputeType_t k = [] { + if (const char* e = std::getenv("VT_ROCM_GEMM_COMPUTE")) { + if (std::strcmp(e, "32f") == 0) return HIPBLAS_COMPUTE_32F; + if (std::strcmp(e, "16f") == 0) return HIPBLAS_COMPUTE_16F; + if (std::strcmp(e, "16bf") == 0) return HIPBLAS_COMPUTE_32F_FAST_16BF; + } + return HIPBLAS_COMPUTE_32F; + }(); + return k; +} + +// VT_ROCM_GEMV=1 enables naive M=1 BF16 GEMV (default OFF — hipblas faster on gfx1201). +bool GemvEnabled() { + static const bool on = [] { + if (const char* e = std::getenv("VT_ROCM_GEMV")) return e[0] == '1'; + return false; + }(); + return on; +} + +// y[n] = alpha * dot(x[0:K], W[n,0:K]) + beta * y[n] +// W row-major [N,K]. +// Strategy: cache x in LDS; each block owns a tile of output rows. +// VT_ROCM_GEMV=1 to enable (default off — A/B vs hipblas). +__global__ void Bf16GemvBTKernel(__hip_bfloat16* __restrict__ y, + const __hip_bfloat16* __restrict__ x, + const __hip_bfloat16* __restrict__ W, int N, int K, + float alpha, float beta) { + // Dynamic shared: x_cache[K] floats + extern __shared__ float smem[]; + float* x_cache = smem; + + // Cooperative load of x into LDS + for (int k = static_cast(threadIdx.x); k < K; k += static_cast(blockDim.x)) { + x_cache[k] = __bfloat162float(x[k]); + } + __syncthreads(); + + // Each thread owns one or more output rows + for (int n = static_cast(blockIdx.x * blockDim.x + threadIdx.x); n < N; + n += static_cast(gridDim.x * blockDim.x)) { + const __hip_bfloat16* wrow = W + static_cast(n) * static_cast(K); + float acc = 0.f; + // Unroll-friendly scalar loop; x from LDS + int k = 0; + for (; k + 3 < K; k += 4) { + acc += x_cache[k] * __bfloat162float(wrow[k]); + acc += x_cache[k + 1] * __bfloat162float(wrow[k + 1]); + acc += x_cache[k + 2] * __bfloat162float(wrow[k + 2]); + acc += x_cache[k + 3] * __bfloat162float(wrow[k + 3]); + } + for (; k < K; ++k) acc += x_cache[k] * __bfloat162float(wrow[k]); + float v = alpha * acc; + if (beta != 0.f) v += beta * __bfloat162float(y[n]); + y[n] = __float2bfloat16(v); + } +} + +void Bf16GemvBT(hipStream_t s, void* out, const void* a, const void* b, int N, int K, + float alpha, float beta) { + constexpr int kBlock = 256; + if (K <= 0 || N <= 0) return; + const int grid = (N + kBlock - 1) / kBlock; + const size_t shmem = static_cast(K) * sizeof(float); + if (shmem > 48 * 1024) return; // caller falls through if we no-op — keep K small + Bf16GemvBTKernel<<>>( + static_cast<__hip_bfloat16*>(out), static_cast(a), + static_cast(b), N, K, alpha, beta); +} + +// VT_ROCM_HIPBLASLT=1 enables. Default OFF — heuristic path aborted on gfx1201 in lab. +bool LtEnabled() { + static const bool on = [] { + if (const char* e = std::getenv("VT_ROCM_HIPBLASLT")) return e[0] == '1'; + return false; + }(); + return on; +} + +struct LtHandle { + hipblasLtHandle_t h = nullptr; + LtHandle() { + if (hipblasLtCreate(&h) != HIPBLAS_STATUS_SUCCESS) h = nullptr; + } + ~LtHandle() { + if (h) (void)hipblasLtDestroy(h); + } +}; + +LtHandle& GetLt() { + static thread_local LtHandle lt; + return lt; +} + +struct LtAlgoCache { + hipblasLtMatmulAlgo_t algo{}; + size_t workspace = 0; + bool ok = false; +}; + +// key: M,N,K,alpha_beta_flag (0=ab1.0/0.0, 1=general) +using LtKey = std::tuple; +struct LtKeyHash { + size_t operator()(const LtKey& k) const { + auto [a, b, c, d] = k; + return (static_cast(a) * 1315423911u) ^ (static_cast(b) * 2654435761u) ^ + (static_cast(c) * 97531u) ^ static_cast(d); + } +}; + +std::unordered_map& LtCache() { + static thread_local std::unordered_map m; + return m; +} + +void* LtWorkspace(size_t need) { + static thread_local void* ws = nullptr; + static thread_local size_t cap = 0; + if (need == 0) return nullptr; + if (need > cap) { + if (ws) (void)hipFree(ws); + ws = nullptr; + cap = 0; + if (hipMalloc(&ws, need) != hipSuccess) return nullptr; + cap = need; + } + return ws; +} + +// C[M,N] row-major = alpha * A[M,K] @ B[N,K]^T + beta * C +// Same col-major view as hipblasGemmEx(OP_T, OP_N, m=N, n=M, k=K, B, K, A, K, C, N). +bool MatmulBTLt(hipStream_t stream, void* C, const void* A, const void* B, int M, int N, int K, + float alpha, float beta) { + if (!LtEnabled() || M <= 0 || N <= 0 || K <= 0) return false; + auto& lt = GetLt(); + if (!lt.h) return false; + + const int ab_flag = (alpha == 1.f && beta == 0.f) ? 0 : 1; + LtKey key{M, N, K, ab_flag}; + auto& cache = LtCache()[key]; + + hipblasLtMatrixLayout_t layoutA = nullptr, layoutB = nullptr, layoutC = nullptr; + hipblasLtMatmulDesc_t matmulDesc = nullptr; + hipblasLtMatmulPreference_t pref = nullptr; + + auto cleanup = [&]() { + if (layoutA) (void)hipblasLtMatrixLayoutDestroy(layoutA); + if (layoutB) (void)hipblasLtMatrixLayoutDestroy(layoutB); + if (layoutC) (void)hipblasLtMatrixLayoutDestroy(layoutC); + if (matmulDesc) (void)hipblasLtMatmulDescDestroy(matmulDesc); + if (pref) (void)hipblasLtMatmulPreferenceDestroy(pref); + }; + + // Col-major views of row-major buffers (default ORDER_COL): + // weight B_rm[N,K] -> A_cm is K x N, ld=K + // act A_rm[M,K] -> B_cm is K x M, ld=K + // out C_rm[M,N] -> C_cm is N x M, ld=N + if (hipblasLtMatrixLayoutCreate(&layoutA, HIP_R_16BF, K, N, K) != HIPBLAS_STATUS_SUCCESS) { + cleanup(); + return false; + } + if (hipblasLtMatrixLayoutCreate(&layoutB, HIP_R_16BF, K, M, K) != HIPBLAS_STATUS_SUCCESS) { + cleanup(); + return false; + } + if (hipblasLtMatrixLayoutCreate(&layoutC, HIP_R_16BF, N, M, N) != HIPBLAS_STATUS_SUCCESS) { + cleanup(); + return false; + } + + if (hipblasLtMatmulDescCreate(&matmulDesc, HIPBLAS_COMPUTE_32F, HIP_R_32F) != + HIPBLAS_STATUS_SUCCESS) { + cleanup(); + return false; + } + // C_cm = A_cm^T * B_cm => (N x K)*(K x M) = N x M == row C[M,N] + const hipblasOperation_t opA = HIPBLAS_OP_T; + const hipblasOperation_t opB = HIPBLAS_OP_N; + (void)hipblasLtMatmulDescSetAttribute(matmulDesc, HIPBLASLT_MATMUL_DESC_TRANSA, &opA, + sizeof(opA)); + (void)hipblasLtMatmulDescSetAttribute(matmulDesc, HIPBLASLT_MATMUL_DESC_TRANSB, &opB, + sizeof(opB)); + + if (!cache.ok) { + if (hipblasLtMatmulPreferenceCreate(&pref) != HIPBLAS_STATUS_SUCCESS) { + cleanup(); + return false; + } + const uint64_t max_ws = 64ull << 20; + (void)hipblasLtMatmulPreferenceSetAttribute(pref, HIPBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, + &max_ws, sizeof(max_ws)); + hipblasLtMatmulHeuristicResult_t heur{}; + int returned = 0; + const hipblasStatus_t hs = hipblasLtMatmulAlgoGetHeuristic( + lt.h, matmulDesc, layoutA, layoutB, layoutC, layoutC, pref, 1, &heur, &returned); + if (hs != HIPBLAS_STATUS_SUCCESS || returned <= 0 || heur.state != HIPBLAS_STATUS_SUCCESS) { + cleanup(); + static thread_local int fail_log = 0; + if (fail_log++ < 3) { + std::fprintf(stderr, "hipBLASLt COL heuristic fail M=%d N=%d K=%d st=%d ret=%d\n", M, N, + K, static_cast(hs), returned); + } + LtCache()[key] = LtAlgoCache{}; + return false; + } + cache.algo = heur.algo; + cache.workspace = heur.workspaceSize; + cache.ok = true; + std::fprintf(stderr, "hipBLASLt COL ok M=%d N=%d K=%d ws=%zu\n", M, N, K, cache.workspace); + } + + void* ws = LtWorkspace(cache.workspace); + if (cache.workspace > 0 && !ws) { + cleanup(); + return false; + } + + // Pass weight as A, act as B (matches GemmEx argument order for OP_T,OP_N) + const hipblasStatus_t st = + hipblasLtMatmul(lt.h, matmulDesc, &alpha, B, layoutA, A, layoutB, &beta, C, layoutC, C, + layoutC, &cache.algo, ws, cache.workspace, stream); + cleanup(); + if (st != HIPBLAS_STATUS_SUCCESS) { + cache.ok = false; + std::fprintf(stderr, "hipBLASLt COL matmul fail st=%d M=%d N=%d K=%d\n", static_cast(st), + M, N, K); + return false; + } + return true; +} + } // namespace // out[M,N] = a[M,K] @ b[K,N] (row-major contiguous / a may have row stride) @@ -133,7 +386,7 @@ void MatmulKernelRocm(Queue& q, Tensor& out, const Tensor& a, const Tensor& b) { /*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, + /*ldc=*/static_cast(N), GemmCompute(a.dtype), HIPBLAS_GEMM_DEFAULT), "hipblasGemmEx NN"); } @@ -166,6 +419,19 @@ void MatmulBTKernelRocm(Queue& q, Tensor& out, const Tensor& a, const Tensor& b) throw std::runtime_error("vt rocm: matmul_bt: bad a stride"); } + // Decode: M=1 BF16 GEMV + if (M == 1 && bf16 && out.dtype == DType::kBF16 && a.stride[0] == K && GemvEnabled()) { + Bf16GemvBT(s, out.data, a.data, b.data, static_cast(N), static_cast(K), 1.f, 0.f); + return; + } + + // hipBLASLt (BF16 contiguous) + if (bf16 && out.dtype == DType::kBF16 && a.stride[0] == K && + MatmulBTLt(s, out.data, a.data, b.data, static_cast(M), static_cast(N), + static_cast(K), 1.f, 0.f)) { + return; + } + auto ctx = GetBlas(q.device.index, s); const float alpha = 1.f, beta = 0.f; const hipDataType at = ToBlasType(a.dtype); @@ -178,9 +444,195 @@ void MatmulBTKernelRocm(Queue& q, Tensor& out, const Tensor& a, const Tensor& b) /*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, + /*ldc=*/static_cast(N), GemmCompute(a.dtype), HIPBLAS_GEMM_DEFAULT), "hipblasGemmEx BT"); } +// out = alpha * (a @ b^T) + beta * out +void MatmulBTAlphaBetaRocm(Queue& q, void* out, const void* a, const void* b, int M, int N, + int K, float alpha, float beta, DType dtype) { + if (M == 0 || N == 0) return; + if (dtype != DType::kBF16 && dtype != DType::kF32) { + throw std::runtime_error("vt rocm: matmul_bt_ab: dtype"); + } + hipStream_t s = static_cast(q.handle); + if (K == 0) { + if (beta == 0.f) { + CheckHip(hipMemsetAsync(out, 0, + static_cast(M) * static_cast(N) * SizeOf(dtype), + s), + "bt_ab k0"); + } + return; + } + // Decode path: custom BF16 GEMV often beats hipblasGemmEx at M=1. + if (M == 1 && dtype == DType::kBF16 && GemvEnabled()) { + Bf16GemvBT(s, out, a, b, N, K, alpha, beta); + return; + } + if (dtype == DType::kBF16 && MatmulBTLt(s, out, a, b, M, N, K, alpha, beta)) { + return; + } + auto ctx = GetBlas(q.device.index, s); + const hipDataType at = ToBlasType(dtype); + CheckBlas(hipblasGemmEx(ctx.handle, HIPBLAS_OP_T, HIPBLAS_OP_N, + /*m=*/N, /*n=*/M, /*k=*/K, &alpha, b, at, /*ldb=*/K, a, at, + /*lda=*/K, &beta, out, at, /*ldc=*/N, GemmCompute(dtype), + HIPBLAS_GEMM_DEFAULT), + "hipblasGemmEx BT alpha/beta"); +} + +// General strided-batch BT: a may also be batched (strideA elements between batches). +// out[g,M,N] = a_g[M,K] @ b_g[N,K]^T +void MatmulBTStridedBatchFullKernelRocm(Queue& q, void* out, const void* a, const void* b, + int batch, int M, int N, int K, long long strideA, + DType dtype) { + if (batch <= 0 || M == 0 || N == 0) return; + if (dtype != DType::kBF16 && dtype != DType::kF32) { + throw std::runtime_error("vt rocm: matmul_bt_batch: dtype"); + } + hipStream_t s = static_cast(q.handle); + if (K == 0) { + CheckHip(hipMemsetAsync(out, 0, + static_cast(batch) * static_cast(M) * + static_cast(N) * SizeOf(dtype), + s), + "bt_batch k0"); + return; + } + auto ctx = GetBlas(q.device.index, s); + const float alpha = 1.f, beta = 0.f; + const hipDataType at = ToBlasType(dtype); + const long long strideB = static_cast(N) * K; + const long long strideC = static_cast(M) * N; + CheckBlas(hipblasGemmStridedBatchedEx( + ctx.handle, HIPBLAS_OP_T, HIPBLAS_OP_N, + /*m=*/N, /*n=*/M, /*k=*/K, &alpha, b, at, /*ldb=*/K, strideB, a, at, + /*lda=*/K, strideA, &beta, out, at, /*ldc=*/N, strideC, batch, + GemmCompute(dtype), HIPBLAS_GEMM_DEFAULT), + "hipblasGemmStridedBatchedEx BT"); +} + +// Batched BT via pointer arrays (no weight gather): out[g] = a @ b[g]^T +// a shared; b_ptrs[g] -> [N,K]; out_ptrs[g] -> [M,N] +void MatmulBTPointerBatchKernelRocm(Queue& q, void** out_ptrs, const void* a, + void** b_ptrs, int batch, int M, int N, int K, + DType dtype) { + if (batch <= 0 || M == 0 || N == 0) return; + if (dtype != DType::kBF16 && dtype != DType::kF32) { + throw std::runtime_error("vt rocm: matmul_bt_ptr_batch: dtype"); + } + hipStream_t s = static_cast(q.handle); + if (K == 0) { + for (int g = 0; g < batch; ++g) + CheckHip(hipMemsetAsync(out_ptrs[g], 0, + static_cast(M) * static_cast(N) * SizeOf(dtype), + s), + "bt_ptr k0"); + return; + } + auto ctx = GetBlas(q.device.index, s); + const float alpha = 1.f, beta = 0.f; + const hipDataType at = ToBlasType(dtype); + + // Reuse device pointer tables (per-device, grows to max batch seen). + struct PtrTables { + void* A = nullptr; + void* B = nullptr; + void* C = nullptr; + int cap = 0; + }; + static std::mutex mu; + static std::unordered_map tabs; + PtrTables* tab = nullptr; + { + std::lock_guard lock(mu); + tab = &tabs[q.device.index]; + if (tab->cap < batch) { + if (tab->A) CheckHip(hipFree(tab->A), "freeA"); + if (tab->B) CheckHip(hipFree(tab->B), "freeB"); + if (tab->C) CheckHip(hipFree(tab->C), "freeC"); + const size_t psz = static_cast(batch) * sizeof(void*); + CheckHip(hipMalloc(&tab->A, psz), "ptrA"); + CheckHip(hipMalloc(&tab->B, psz), "ptrB"); + CheckHip(hipMalloc(&tab->C, psz), "ptrC"); + tab->cap = batch; + } + } + std::vector a_host(static_cast(batch), a); + const size_t psz = static_cast(batch) * sizeof(void*); + CheckHip(hipMemcpyAsync(tab->A, a_host.data(), psz, hipMemcpyHostToDevice, s), "H2D A"); + CheckHip(hipMemcpyAsync(tab->B, b_ptrs, psz, hipMemcpyHostToDevice, s), "H2D B"); + CheckHip(hipMemcpyAsync(tab->C, out_ptrs, psz, hipMemcpyHostToDevice, s), "H2D C"); + + CheckBlas(hipblasGemmBatchedEx( + ctx.handle, HIPBLAS_OP_T, HIPBLAS_OP_N, + /*m=*/N, /*n=*/M, /*k=*/K, &alpha, + reinterpret_cast(tab->B), at, /*ldb=*/K, + reinterpret_cast(tab->A), at, /*lda=*/K, &beta, + reinterpret_cast(tab->C), at, /*ldc=*/N, batch, GemmCompute(dtype), + HIPBLAS_GEMM_DEFAULT), + "hipblasGemmBatchedEx BT"); +} + +// Pointer-array BT with per-batch A and B: out[g] = a[g] @ b[g]^T +void MatmulBTPointerBatchABKernelRocm(Queue& q, void** out_ptrs, void** a_ptrs, + void** b_ptrs, int batch, int M, int N, int K, + DType dtype) { + if (batch <= 0 || M == 0 || N == 0) return; + if (dtype != DType::kBF16 && dtype != DType::kF32) { + throw std::runtime_error("vt rocm: matmul_bt_ptr_ab: dtype"); + } + hipStream_t s = static_cast(q.handle); + if (K == 0) { + for (int g = 0; g < batch; ++g) + CheckHip(hipMemsetAsync(out_ptrs[g], 0, + static_cast(M) * static_cast(N) * SizeOf(dtype), + s), + "bt_ptr_ab k0"); + return; + } + auto ctx = GetBlas(q.device.index, s); + const float alpha = 1.f, beta = 0.f; + const hipDataType at = ToBlasType(dtype); + + struct PtrTables { + void* A = nullptr; + void* B = nullptr; + void* C = nullptr; + int cap = 0; + }; + static std::mutex mu; + static std::unordered_map tabs; + PtrTables* tab = nullptr; + { + std::lock_guard lock(mu); + tab = &tabs[q.device.index]; + if (tab->cap < batch) { + if (tab->A) CheckHip(hipFree(tab->A), "freeA"); + if (tab->B) CheckHip(hipFree(tab->B), "freeB"); + if (tab->C) CheckHip(hipFree(tab->C), "freeC"); + const size_t psz = static_cast(batch) * sizeof(void*); + CheckHip(hipMalloc(&tab->A, psz), "ptrA"); + CheckHip(hipMalloc(&tab->B, psz), "ptrB"); + CheckHip(hipMalloc(&tab->C, psz), "ptrC"); + tab->cap = batch; + } + } + const size_t psz = static_cast(batch) * sizeof(void*); + CheckHip(hipMemcpyAsync(tab->A, a_ptrs, psz, hipMemcpyHostToDevice, s), "H2D A"); + CheckHip(hipMemcpyAsync(tab->B, b_ptrs, psz, hipMemcpyHostToDevice, s), "H2D B"); + CheckHip(hipMemcpyAsync(tab->C, out_ptrs, psz, hipMemcpyHostToDevice, s), "H2D C"); + + CheckBlas(hipblasGemmBatchedEx( + ctx.handle, HIPBLAS_OP_T, HIPBLAS_OP_N, + /*m=*/N, /*n=*/M, /*k=*/K, &alpha, + reinterpret_cast(tab->B), at, /*ldb=*/K, + reinterpret_cast(tab->A), at, /*lda=*/K, &beta, + reinterpret_cast(tab->C), at, /*ldc=*/N, batch, GemmCompute(dtype), + HIPBLAS_GEMM_DEFAULT), + "hipblasGemmBatchedEx BT AB"); +} + } // namespace vt::rocm diff --git a/src/vt/rocm/rocm_moe_router.hip b/src/vt/rocm/rocm_moe_router.hip new file mode 100644 index 000000000..d9461e792 --- /dev/null +++ b/src/vt/rocm/rocm_moe_router.hip @@ -0,0 +1,134 @@ +// ROCm ungrouped MoE router: softmax + greedy top-k (lowest-index tie-break). +#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 moe_router: ") + what + ": " + + hipGetErrorString(err)); + } +} + +// One block per token. E <= 512, K <= 64. Shared: scores[E] + mask work in registers. +__global__ void MoeRouterTopKUngroupedKernel(const float* __restrict__ logits, + float* __restrict__ weights, + int32_t* __restrict__ indices, int T, int E, + int K, int renormalize, float scale) { + const int t = static_cast(blockIdx.x); + if (t >= T) return; + extern __shared__ float smem[]; + float* scores = smem; // [E] + + const float* row = logits + static_cast(t) * static_cast(E); + // Load + max + float mx = -INFINITY; + for (int j = static_cast(threadIdx.x); j < E; j += static_cast(blockDim.x)) { + const float v = row[j]; + scores[j] = v; + mx = fmaxf(mx, v); + } + // block reduce max + __shared__ float red[256]; + red[threadIdx.x] = mx; + __syncthreads(); + for (int s = static_cast(blockDim.x) / 2; s > 0; s >>= 1) { + if (static_cast(threadIdx.x) < s) red[threadIdx.x] = fmaxf(red[threadIdx.x], red[threadIdx.x + s]); + __syncthreads(); + } + mx = red[0]; + + // softmax + float sum = 0.f; + for (int j = static_cast(threadIdx.x); j < E; j += static_cast(blockDim.x)) { + const float e = expf(scores[j] - mx); + scores[j] = e; + sum += e; + } + red[threadIdx.x] = sum; + __syncthreads(); + for (int s = static_cast(blockDim.x) / 2; s > 0; s >>= 1) { + if (static_cast(threadIdx.x) < s) red[threadIdx.x] += red[threadIdx.x + s]; + __syncthreads(); + } + sum = red[0]; + const float inv = (sum > 0.f) ? (1.f / sum) : 0.f; + for (int j = static_cast(threadIdx.x); j < E; j += static_cast(blockDim.x)) { + float p = scores[j] * inv; + if (!isfinite(p)) p = 0.f; + scores[j] = p; + } + __syncthreads(); + + // Greedy top-k with strict > (lowest index wins ties). Thread 0 only — E small. + if (threadIdx.x == 0) { + float denom = 0.f; + float* wout = weights + static_cast(t) * static_cast(K); + int32_t* iout = indices + static_cast(t) * static_cast(K); + for (int ki = 0; ki < K; ++ki) { + int best = 0; + float best_v = -INFINITY; + for (int j = 0; j < E; ++j) { + const float v = scores[j]; + if (v > best_v) { + best_v = v; + best = j; + } + } + const float w = scores[best]; + scores[best] = -INFINITY; + wout[ki] = w; + iout[ki] = best; + denom += w; + } + if (renormalize) { + if (!(denom > 0.f)) denom = 1.f; + for (int ki = 0; ki < K; ++ki) wout[ki] /= denom; + } + if (scale != 1.f) { + for (int ki = 0; ki < K; ++ki) wout[ki] *= scale; + } + } +} + +} // namespace + +void MoeRouterTopKKernelRocm(Queue& q, Tensor& weights, Tensor& indices, const Tensor& logits, + const MoeRouterTopKArgs& args, const Tensor* bias) { + // Grouped / bias path: fall back via host (rare for Gemma4). + if (args.num_expert_group > 0 || bias != nullptr || + args.scoring_func != MoeScoringFunc::kSoftmax) { + throw std::runtime_error( + "vt rocm: MoeRouterTopK only implements ungrouped softmax without bias; " + "use CPU backend for grouped/bias routers"); + } + if (logits.dtype != DType::kF32 || weights.dtype != DType::kF32 || + indices.dtype != DType::kI32) { + throw std::runtime_error("vt rocm: MoeRouterTopK dtypes must be f32/f32/i32"); + } + const int T = static_cast(logits.shape[0]); + const int E = static_cast(logits.shape[1]); + const int K = args.top_k; + if (T == 0 || K == 0) return; + if (E > 512 || K > 64) { + throw std::runtime_error("vt rocm: MoeRouterTopK E>512 or K>64 not supported"); + } + hipStream_t s = static_cast(q.handle); + const int threads = 256; + const size_t shmem = static_cast(E) * sizeof(float); + hipLaunchKernelGGL(MoeRouterTopKUngroupedKernel, dim3(T), dim3(threads), shmem, s, + static_cast(logits.data), static_cast(weights.data), + static_cast(indices.data), T, E, K, args.renormalize ? 1 : 0, + args.routed_scaling_factor); + CheckHip(hipGetLastError(), "MoeRouterTopK launch"); +} + +} // namespace vt::rocm diff --git a/src/vt/rocm/rocm_ops.hip b/src/vt/rocm/rocm_ops.hip index 5667cf00f..bac34a990 100644 --- a/src/vt/rocm/rocm_ops.hip +++ b/src/vt/rocm/rocm_ops.hip @@ -36,6 +36,8 @@ void PagedAttentionKernelRocm(Queue& q, Tensor& out, const Tensor& query, const const PagedAttentionArgs& args); void GeluTanhKernelRocm(Queue& q, Tensor& out, const Tensor& x); void GeluErfKernelRocm(Queue& q, Tensor& out, const Tensor& x); +void MoeRouterTopKKernelRocm(Queue& q, Tensor& weights, Tensor& indices, const Tensor& logits, + const MoeRouterTopKArgs& args, const Tensor* bias); namespace { @@ -90,6 +92,9 @@ struct Registrar { RegisterOp(OpId::kPagedAttention, DeviceType::kROCM, reinterpret_cast( static_cast(&PagedAttentionKernelRocm))); + RegisterOp(OpId::kMoeRouterTopK, DeviceType::kROCM, + reinterpret_cast( + static_cast(&MoeRouterTopKKernelRocm))); } } registrar; diff --git a/src/vt/rocm/rocm_rmsnorm.hip b/src/vt/rocm/rocm_rmsnorm.hip index 24cc3e815..2792e74b8 100644 --- a/src/vt/rocm/rocm_rmsnorm.hip +++ b/src/vt/rocm/rocm_rmsnorm.hip @@ -129,6 +129,34 @@ void LaunchRmsNorm(hipStream_t s, Tensor& out, const Tensor& x, const Tensor& w, } } +// out = rmsnorm(x, w) + addend — Gemma-4 residual join (NOT residual+=x then norm). +template +__global__ void RmsNormPlusAddKernel(T* out, const T* x, const T* w, const T* add, int64_t h, + float eps, bool gemma) { + const int64_t row = blockIdx.x; + const T* xrow = x + row * h; + const T* arow = add + row * h; + T* orow = out + row * h; + __shared__ float partial[kBlock]; + float acc = 0.0f; + for (int64_t j = threadIdx.x; j < h; j += kBlock) { + const float v = Load(xrow, j); + acc += v * v; + } + partial[threadIdx.x] = acc; + __syncthreads(); + for (int s = kBlock / 2; s > 0; s /= 2) { + if (static_cast(threadIdx.x) < s) partial[threadIdx.x] += partial[threadIdx.x + s]; + __syncthreads(); + } + const float inv = 1.0f / sqrtf(partial[0] / static_cast(h) + eps); + for (int64_t j = threadIdx.x; j < h; j += kBlock) { + float wj = Load(w, j); + if (gemma) wj += 1.0f; + Store(orow, j, Load(xrow, j) * inv * wj + Load(arow, j)); + } +} + } // namespace // The registered RmsNormFn (include/vt/ops.h:810). Signature is the shared vt @@ -149,4 +177,121 @@ void RmsNormKernelRocm(Queue& q, Tensor& out, const Tensor& x, const Tensor& w, } } +// out = rmsnorm(x, w) + addend +void RmsNormPlusAddRocm(Queue& q, Tensor& out, const Tensor& x, const Tensor& w, + const Tensor& addend, const RmsNormArgs& args) { + VT_CHECK(w.dtype == x.dtype && out.dtype == x.dtype && addend.dtype == x.dtype, + "rocm rmsnorm+add: dtype match"); + VT_CHECK(x.rank >= 2 && out.shape[0] == x.shape[0] && out.shape[1] == x.shape[1] && + addend.shape[0] == x.shape[0] && addend.shape[1] == x.shape[1], + "rocm rmsnorm+add: shape"); + const int64_t t = x.shape[0], h = x.shape[1]; + if (t == 0 || h == 0) return; + hipStream_t s = static_cast(q.handle); + if (x.dtype == DType::kBF16) { + RmsNormPlusAddKernel<__hip_bfloat16><<(t), kBlock, 0, s>>>( + out.Ptr<__hip_bfloat16>(), x.Ptr<__hip_bfloat16>(), w.Ptr<__hip_bfloat16>(), + addend.Ptr<__hip_bfloat16>(), h, args.eps, args.gemma); + } else if (x.dtype == DType::kF32) { + RmsNormPlusAddKernel<<(t), kBlock, 0, s>>>( + out.Ptr(), x.Ptr(), w.Ptr(), addend.Ptr(), h, args.eps, + args.gemma); + } else { + VT_CHECK(false, "rocm rmsnorm+add dtype"); + } +} + +// out = rmsnorm(rmsnorm(x1,w1) + rmsnorm(x2,w2), w3) + residual +// Gemma-4 MoE post-FF: dual expert streams + join residual. +template +__global__ void DualRmsNormPlusResKernel(T* out, const T* x1, const T* w1, const T* x2, + const T* w2, const T* w3, const T* residual, int64_t h, + float eps1, float eps2, float eps3, bool gemma) { + const int64_t row = blockIdx.x; + const T* x1r = x1 + row * h; + const T* x2r = x2 + row * h; + const T* rr = residual + row * h; + T* orow = out + row * h; + __shared__ float p1[kBlock]; + __shared__ float p2[kBlock]; + __shared__ float p3[kBlock]; + + float a1 = 0.f, a2 = 0.f; + for (int64_t j = threadIdx.x; j < h; j += kBlock) { + const float v1 = Load(x1r, j); + const float v2 = Load(x2r, j); + a1 += v1 * v1; + a2 += v2 * v2; + } + p1[threadIdx.x] = a1; + p2[threadIdx.x] = a2; + __syncthreads(); + for (int s = kBlock / 2; s > 0; s /= 2) { + if (static_cast(threadIdx.x) < s) { + p1[threadIdx.x] += p1[threadIdx.x + s]; + p2[threadIdx.x] += p2[threadIdx.x + s]; + } + __syncthreads(); + } + const float inv1 = 1.f / sqrtf(p1[0] / static_cast(h) + eps1); + const float inv2 = 1.f / sqrtf(p2[0] / static_cast(h) + eps2); + + // Second pass: build mid = n1+n2 and accumulate variance of mid. + // Use shared? too large for H=2816. Recompute n1,n2 per thread. + float a3 = 0.f; + for (int64_t j = threadIdx.x; j < h; j += kBlock) { + float w1j = Load(w1, j); + float w2j = Load(w2, j); + if (gemma) { + w1j += 1.f; + w2j += 1.f; + } + const float mid = Load(x1r, j) * inv1 * w1j + Load(x2r, j) * inv2 * w2j; + a3 += mid * mid; + } + p3[threadIdx.x] = a3; + __syncthreads(); + for (int s = kBlock / 2; s > 0; s /= 2) { + if (static_cast(threadIdx.x) < s) p3[threadIdx.x] += p3[threadIdx.x + s]; + __syncthreads(); + } + const float inv3 = 1.f / sqrtf(p3[0] / static_cast(h) + eps3); + for (int64_t j = threadIdx.x; j < h; j += kBlock) { + float w1j = Load(w1, j); + float w2j = Load(w2, j); + float w3j = Load(w3, j); + if (gemma) { + w1j += 1.f; + w2j += 1.f; + w3j += 1.f; + } + const float mid = Load(x1r, j) * inv1 * w1j + Load(x2r, j) * inv2 * w2j; + Store(orow, j, mid * inv3 * w3j + Load(rr, j)); + } +} + +void DualRmsNormPlusResRocm(Queue& q, Tensor& out, const Tensor& x1, const Tensor& w1, + const Tensor& x2, const Tensor& w2, const Tensor& w3, + const Tensor& residual, const RmsNormArgs& args) { + VT_CHECK(out.dtype == x1.dtype && x1.dtype == x2.dtype && x1.dtype == residual.dtype, + "dual rms dtype"); + const int64_t t = x1.shape[0], h = x1.shape[1]; + if (t == 0 || h == 0) return; + hipStream_t s = static_cast(q.handle); + // plain RMS (no gemma +1) for Gemma-4 layer norms + const bool gemma = args.gemma; + if (x1.dtype == DType::kBF16) { + DualRmsNormPlusResKernel<__hip_bfloat16><<(t), kBlock, 0, s>>>( + out.Ptr<__hip_bfloat16>(), x1.Ptr<__hip_bfloat16>(), w1.Ptr<__hip_bfloat16>(), + x2.Ptr<__hip_bfloat16>(), w2.Ptr<__hip_bfloat16>(), w3.Ptr<__hip_bfloat16>(), + residual.Ptr<__hip_bfloat16>(), h, args.eps, args.eps, args.eps, gemma); + } else if (x1.dtype == DType::kF32) { + DualRmsNormPlusResKernel<<(t), kBlock, 0, s>>>( + out.Ptr(), x1.Ptr(), w1.Ptr(), x2.Ptr(), w2.Ptr(), + w3.Ptr(), residual.Ptr(), h, args.eps, args.eps, args.eps, gemma); + } else { + VT_CHECK(false, "dual rms dtype"); + } +} + } // namespace vt::rocm