diff --git a/edge0-strata/CMakeLists.txt b/edge0-strata/CMakeLists.txt index afc05c5f..7ace4454 100644 --- a/edge0-strata/CMakeLists.txt +++ b/edge0-strata/CMakeLists.txt @@ -2495,6 +2495,18 @@ if(STRATA_NATIVE_EXPERTS) src/prefill/moe_f16_kq.cu src/prefill/ggml_cuda_host.cu "${_strata_quantize_cu}") + # amd-hip port: the latent A stage sources of the CUDA branch. On HIP the lat_q3k_* entry points are stubs that + # exit ("CUDA-only"), so run with STRATA_LATENT_FUSED=0 STRATA_LATENT_FUSED_DU=0 STRATA_LAT_FUSED_Q3K=0; they are + # compiled here so prefill.cpp links. + set(_strata_lat_hip_sources + src/prefill/lat_q3k_down.cu + src/prefill/lat_q3k_gu.cu + src/prefill/lat_q3k_swiglu.cu + src/prefill/lat_q3k_fused.cu + src/prefill/moe_fused_q3k.cu + src/prefill/lat_dense_w11.cu # the latent D / U projections on the gfx11 matrix cores + src/prefill/lat_q3k_w11.cu) # the latent A stage (Q3_K gate/up, IQ4_NL down) on the gfx11 matrix cores + list(APPEND _strata_mmq_sources ${_strata_lat_hip_sources}) if(STRATA_MMQ_KQUANTS) # the same K-quant instances as the CUDA branch: on HIP they serve the dense GGUF projections # of Gemm::native (STRATA_DENSE_MMQ=1), the mixed-quant packs' attention/PLE tensors @@ -2512,6 +2524,8 @@ if(STRATA_NATIVE_EXPERTS) # moe_fused.cu is Strata's own CUDA-shaped source: it takes the HIP compat headers, ggml's MMQ sources do not set_source_files_properties(src/prefill/moe_fused.cu src/prefill/moe_fused_iq.cu src/prefill/moe_f16_kq.cu PROPERTIES INCLUDE_DIRECTORIES "${STRATA_HIP_COMPAT_INCLUDE_DIR}") + set_source_files_properties(${_strata_lat_hip_sources} + PROPERTIES INCLUDE_DIRECTORIES "${STRATA_HIP_COMPAT_INCLUDE_DIR}") # amd-hip port add_library(strata_mmq STATIC ${_strata_mmq_sources}) target_include_directories(strata_mmq PRIVATE "${_strata_gg}/include" "${_strata_gg}/src" "${_strata_mmq}" @@ -2612,6 +2626,9 @@ if(STRATA_NATIVE_EXPERTS) target_include_directories(latent_cpu_verify PRIVATE ${CMAKE_SOURCE_DIR}/include) target_link_libraries(latent_cpu_verify PRIVATE strata_engine strata_kernels_cpu CUDA::cudart) endif() + if(STRATA_ENABLE_HIP) # amd-hip port: the CUDA block above adds these to the engine; HIP needs them too + target_sources(strata_engine PRIVATE src/core/latent_dense.cpp src/latent/latent_moe.cpp) + endif() if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP) add_executable(native_expert_parity src/kernels/native_expert_parity.cpp) target_include_directories(native_expert_parity PRIVATE ${_strata_gpu_include_directories}) diff --git a/edge0-strata/EDGE0.md b/edge0-strata/EDGE0.md index 4ec3743e..69eca788 100644 --- a/edge0-strata/EDGE0.md +++ b/edge0-strata/EDGE0.md @@ -193,6 +193,132 @@ interpreter. To run `build\strata.exe` by hand, put the CUDA DLLs on `PATH` firs (`set "PATH=%CUDA_PATH%\bin;%PATH%"`). The engine writes temporary files under `%TMP%`, which must be writable for the user who runs `edge0-infer`. +## Build (Windows, AMD, experimental) + +Edge0 145B (Q3_K) has been run on one AMD card with the HIP backend this fork inherits from Strata: a Radeon RX 7900 +XTX (gfx1100, 24 GB) under Windows 11, ROCm 10.2.0a20260930. Other Radeon cards, Linux, Edge0 120B and the NVFP4 +folders have not been tried. `edge0-infer` does not start this build (it expects CUDA), so the server is started by +hand. + +Requirements: Visual Studio 2022 Build Tools with the C++ workload, Python 3.10 or newer, Git, and a current AMD +display driver. No ROCm or HIP SDK installation: the build script installs ROCm from AMD's wheels into `.rocm-win` +(about 7 GB). + +1. Build and package the engine for the card's architecture, in this directory: + + ```bat + set STRATA_HIP_ARCHS=gfx1100 + set STRATA_NO_VISION=1 + tools\hip\build_windows.bat + ``` + + The result is `dist\strata-windows-x64-hip.zip` (`strata.exe`, the ROCm DLLs in `rocm\bin`, the hipBLASLt tuning + table in `tuning\`). Unpack it into a folder, for example `engine-hip\`. +2. Create the server's Python environment as in step 3 of [Build (Windows, NVIDIA)](#build-windows-nvidia). +3. Write a server configuration from `edge0/models/edge0-145b.json`: replace the placeholders by hand, set `exe` to + the unpacked `strata.exe`, add `"backend": "hip"`, and set `lib_dirs` to the unpacked `rocm\bin`. In `env`: + + | Variable | Value | Why | + | --- | --- | --- | + | `STRATA_LATENT_FUSED`, `STRATA_LATENT_FUSED_DU`, `STRATA_LAT_FUSED_Q3K` | `0` | the fused latent kernels are CUDA-only; the engine exits if they are selected on HIP | + | `STRATA_LATENT_MMQ` | `1` (optional) | the latent A stage of the prompt on MMQ; about a third faster than the default path on this card | + | `STRATA_LATENT_DENSE_FMT` | `8` or `4` | the engine reads the number, not `q8` / `q4` | + | `STRATA_LATENT_A`, `STRATA_LATENT_DENSE` | paths with `/` | `labels.bin` is looked up beside `STRATA_LATENT_A` by its last `/`, so a path with only `\` is not found | + | `STRATA_HIPBLASLT_TUNING` | the table in `tuning\` | as in upstream Strata's AMD configurations | + +4. Start it: `.venv\Scripts\python.exe serve\server.py --engine strata --config --port 8000`. + +Measured with `tools\strata_bench.py` on the card above (Ryzen 7 5800X3D, 64 GB RAM, `--vram-reserve-mib 3072`, about +4100 experts resident; medians): + +| A stage of the prompt | Chat decode tok/s | 4K prompt tok/s | 16K prompt tok/s | Agent decode tok/s | +| --- | ---: | ---: | ---: | ---: | +| default (`native_expert_grouped`) | 37.9 | 243.6 | 241.3 | 43.6 | +| `STRATA_LATENT_MMQ=1` | 40.4 | 335.0 | 319.3 | 47.2 | + +The prompt speed is far below the NVIDIA figures because the fused latent kernels have no HIP version yet. Answer +quality on this build has only been checked by hand. + +## Build (Linux, AMD, experimental) + +Edge0 145B (Q3_K) runs on a Radeon RX 7900 XTX (gfx1100, 24 GB) under Linux with ROCm 7.2.4. The expert stage and +the two dense projections of a latent layer have gfx11 kernels (`src/prefill/lat_q3k_w11.cu`, +`src/prefill/lat_dense_w11.cu`). `edge0-infer` does not start this build, so start the server by hand. + +A layer split across two cards works (`"gpu": [0, 1]`). Each card holds the latent D and U weights of its own +layers, and on a split the token embedding table stays in mapped host memory whatever `STRATA_EMBED_VRAM` says. A +card cannot read another card's VRAM: HIP answers "an illegal memory access was encountered", and the window or the +draft layer then never finishes. Add `--resident-experts` to keep in RAM only the experts that no card holds. + +Do not set `--pcie-mode dma` on a split. It was not tested again after the fix above. Before the fix, the first +expert that a card did not hold was never copied, and the request stopped with `verify: timed out at layer N`. + +Requirements: ROCm 7 with the HIP compiler, hipBLAS and hipBLASLt (the commands below were run in the +`docker.io/rocm/dev-ubuntu-24.04:7.2.4-complete` image), CMake 3.24 or newer, Ninja, and Python 3.10 or newer. + +1. Unpack llama.cpp into `third_party/llama.cpp` as in step 1 of [Build (Linux, NVIDIA)](#build-linux-nvidia). +2. Build the engine. `STRATA_PREFILL_MMQ=ON` is required: without it the link stops with undefined symbols such as + `strata::prefill::fused_q3k::experts_latent_q3k`. + + ``` + export HIP_PATH="$(hipconfig -R)" + cmake -S . -B build-hip -G Ninja -DCMAKE_BUILD_TYPE=Release \ + -DSTRATA_ENABLE_HIP=ON -DSTRATA_ENABLE_CUDA=OFF -DCMAKE_HIP_ARCHITECTURES=gfx1100 \ + -DCMAKE_HIP_COMPILER="$(hipconfig -l)/clang++" -DCMAKE_PREFIX_PATH=/opt/rocm \ + -DSTRATA_PREFILL_MMQ=ON -DSTRATA_NATIVE_EXPERTS=ON \ + -DSTRATA_GGML_DIR="$PWD/third_party/llama.cpp" + cmake --build build-hip --target strata + ``` + +3. Copy the hipBLASLt tuning table for the card and the hipBLASLt version into `tuning/` beside the engine. The + engine reads it from there and prints `hipBLASLt tuning enabled`. Without it the dense prompt products use plain + hipBLAS, and a 5.3K prompt ran at 725 tok/s where it runs at 1150 tok/s with the table. ROCm 7.2.4 has hipBLASLt + 1.2.2, which is table `100202`: + + ``` + mkdir -p build-hip/tuning + cp tools/hip/gfx1100-hipblaslt-100202.txt build-hip/tuning/ + ``` + +4. Create the server's Python environment as in step 3 of [Build (Linux, NVIDIA)](#build-linux-nvidia). +5. Write a server configuration from `edge0/models/edge0-145b.json`. Replace the placeholders, set `exe` to + `build-hip/strata`, add `"backend": "hip"`, and set `lib_dirs` to the ROCm library directory (`/opt/rocm/lib`). + Set `cwd` and `log` to a directory that the server can write to. In `env`, set `STRATA_LATENT_FUSED` and + `STRATA_LAT_FUSED_Q3K` to `1`, and `STRATA_LATENT_DENSE_FMT` to `8`. Do not set `--pool-workers` above the number + of physical cores; without the option the engine chooses. +6. Start it: `.venv/bin/python serve/server.py --engine strata --config --port 8000`. The process needs + `/dev/kfd` and `/dev/dri`. With more than one card, select one with `HIP_VISIBLE_DEVICES`. + +Use a card on a full-width PCIe link. The engine moves experts between host memory and the card while it decodes, +and it prints each card's link speed at the start of a layer split (`link 0.8 GB/s`). `current_link_width` of each +bridge above the card in `/sys/bus/pci/devices` shows a narrow link. + +Measured on a Ryzen 5 5600X with 64 GB RAM and two RX 7900 XTX cards, one card in use, `--vram-reserve-mib 700`, +6652 experts resident (the card also held 1.8 GiB of another process). The client sends three chat questions of 150 +tokens and one fresh prompt of each size, thinking off, one run each; it is not `tools/strata_bench.py`. + +| Card link | Configuration | Chat decode tok/s | 5.3K prompt tok/s | 21K prompt tok/s | +| --- | --- | ---: | ---: | ---: | +| PCIe 4.0 x16 | all of the above | not measured | 1181.8 | 1302.1 | +| PCIe 4.0 x16 | expert stage on ggml MMQ: `STRATA_LATENT_FUSED=0`, `STRATA_LAT_FUSED_Q3K=0`, `STRATA_LATENT_MMQ=1` | 46.2 to 51.1 | 1157.2 to 1202.0 | 1285.3 to 1316.1 | +| PCIe 4.0 x16 | MMQ expert stage, no tuning table | not measured | 724.9 | not measured | +| PCIe 4.0 x16 | MMQ expert stage, no tuning table, `STRATA_LATENT_FUSED_DU=0` | 46.4 to 55.1 | 269.1 | 280.2 | +| PCIe 4.0 x16 | native expert stage, no tuning table, `STRATA_LATENT_FUSED_DU=0` | 47.7 to 57.1 | 214.3 | 216.7 | +| PCIe 3.0 x1 | native expert stage, no tuning table, `STRATA_LATENT_FUSED_DU=0` | 17.8 to 23.3 | 108.3 | 129.1 | + +The gfx11 expert stage and the MMQ expert stage ran in different sessions, and the same MMQ configuration measured +1157.2 and 1202.0 tok/s on the 5.3K prompt, so the two stages are level within the run-to-run spread. + +Two cards in a layer split with `--resident-experts`, layers 0 to 21 on the PCIe 3.0 x1 card and 22 to 47 on the x16 +card: the cards held 27.0 GiB of the 37.3 GiB of experts, and the engine kept 21.4 GiB of experts in RAM, where one +card keeps all 37.3 GiB. The process took 32.5 GB of RAM. Chat answers decoded at 32.2 to 51.2 tok/s. A 5.3K prompt +ran at 272.9 to 284.6 tok/s and a 21K prompt at 413.2 tok/s, because the prompt path streams experts to the x1 card. + +Answer quality on this build has only been checked by hand and by the first token's log-probabilities over four +prompts. The gfx11 expert stage moves the top-8 log-probabilities by at most 0.48 against the MMQ stage. The gfx11 +projections move them by at most 0.57 against `STRATA_LATENT_FUSED_DU=0`. The two upstream expert stages (native and +MMQ) differ from each other by up to 2.41 on the same prompts. + ## Launch contract with edge0-infer - `edge0/CONTRACT` holds one integer, the contract version (currently 4). `edge0-infer` refuses to start an diff --git a/edge0-strata/include/strata/core/latent_share.hpp b/edge0-strata/include/strata/core/latent_share.hpp index fa8b7d7c..4849d67e 100644 --- a/edge0-strata/include/strata/core/latent_share.hpp +++ b/edge0-strata/include/strata/core/latent_share.hpp @@ -1,4 +1,4 @@ -// Latent MoE: the one session-level D/U device pair, stored by the engine's LatentTab and pulled by the +// Latent MoE: the session-level D/U device pair of each device, stored by the engine's LatentTab and pulled by the // prefill and verify arms inside their own init (no per-object external wiring to miss). A non-latent pack // never stores a pair, so both accessors stay null. #pragma once @@ -7,8 +7,18 @@ namespace strata { namespace core { -inline uint8_t*& latent_shared_d() { static uint8_t* p = nullptr; return p; } -inline uint8_t*& latent_shared_u() { static uint8_t* p = nullptr; return p; } +// The pair is device memory, so there is one per device: a kernel cannot read another device's buffer (HIP answers +// "an illegal memory access was encountered"). A layer split's later stage pulls the pair of its own device, or +// finds none and uploads its own copy. `device` is the CUDA / HIP device number. +struct LatentSharedPair { + uint8_t *d = nullptr, *u = nullptr; +}; +inline LatentSharedPair& latent_shared(int device) { + static LatentSharedPair pairs[32]; + return pairs[device & 31]; +} +inline uint8_t*& latent_shared_d(int device) { return latent_shared(device).d; } +inline uint8_t*& latent_shared_u(int device) { return latent_shared(device).u; } inline const uint8_t*& latent_xl_current() { static const uint8_t* p = nullptr; return p; } inline const uint32_t*& latent_xlseq_current() { static const uint32_t* p = nullptr; return p; } diff --git a/edge0-strata/include/strata/core/native_head.hpp b/edge0-strata/include/strata/core/native_head.hpp index 29fe9d37..96d59239 100644 --- a/edge0-strata/include/strata/core/native_head.hpp +++ b/edge0-strata/include/strata/core/native_head.hpp @@ -51,7 +51,10 @@ class NativeEmbed { ~NativeEmbed(); NativeEmbed(const NativeEmbed&) = delete; NativeEmbed& operator=(const NativeEmbed&) = delete; - bool load(const std::vector& shards, int64_t n_embd, int64_t n_vocab, std::string& err); + /// `one_device` = false (a layer split): the table stays in mapped host memory whatever STRATA_EMBED_VRAM + /// says, because a later stage's device cannot read another device's VRAM. + bool load(const std::vector& shards, int64_t n_embd, int64_t n_vocab, std::string& err, + bool one_device = true); /// Rows for device token ids. void gather_dev(const int32_t* tokens, int64_t n_tok, float* out, void* stream) const; /// One row for a host token id. diff --git a/edge0-strata/include/strata/hip_compat/cuda_runtime.h b/edge0-strata/include/strata/hip_compat/cuda_runtime.h index 9f2e2f18..2e8cbf3f 100644 --- a/edge0-strata/include/strata/hip_compat/cuda_runtime.h +++ b/edge0-strata/include/strata/hip_compat/cuda_runtime.h @@ -11,6 +11,9 @@ #endif #include #include +#include +#include +#include // Do not let HIP's legacy macro corrupt libstdc++ attribute names. #ifdef __noinline__ #undef __noinline__ @@ -94,17 +97,13 @@ hipError_t mem_get_info(size_t* free_bytes, size_t* total_bytes); #else #define cudaMemGetInfo hipMemGetInfo #endif -#define cudaMemcpy hipMemcpy #define cudaMemcpy2DAsync hipMemcpy2DAsync -#define cudaMemcpyAsync hipMemcpyAsync #define cudaMemcpyPeerAsync hipMemcpyPeerAsync #define cudaMemcpyPeer hipMemcpyPeer #define cudaMemcpyDefault hipMemcpyDefault #define cudaMemcpyDeviceToDevice hipMemcpyDeviceToDevice #define cudaMemcpyDeviceToHost hipMemcpyDeviceToHost #define cudaMemcpyHostToDevice hipMemcpyHostToDevice -#define cudaMemset hipMemset -#define cudaMemsetAsync hipMemsetAsync #define cudaPeekAtLastError hipPeekAtLastError #define cudaRuntimeGetVersion hipRuntimeGetVersion #define cudaSetDevice hipSetDevice @@ -131,6 +130,86 @@ inline hipError_t malloc_host(T** pointer, size_t bytes) { } } // namespace strata::hip_compat +// ---- copies and fills on the buffer's own device. +// CUDA takes a memcpy or a memset of any device's buffer whatever the current device is, and the engine's layer +// split relies on that. HIP runs such a call on the CURRENT device's queue. A blocking hipMemset of another +// device's buffer then returns before that device sees the bytes (measured on two gfx1100 cards, ROCm 7.2.4: a +// kernel launched on the buffer's device right after the call still read the old bytes), so the fill lands late, +// over whatever that device wrote in between. With more than one device, these wrappers make the buffer's device +// current for the call. A call on a stream of its own is left alone: the stream names the device. +// STRATA_HIP_XDEV_LOG=1 prints the first calls that needed the switch. +namespace strata::hip_compat { +inline bool multi_device() { + static const bool multi = [] { + int n = 0; + return hipGetDeviceCount(&n) == hipSuccess && n > 1; + }(); + return multi; +} +// the device that owns `p`, or -1 for host memory and for a pointer HIP does not know +inline int device_of(const void* p) { + hipPointerAttribute_t a{}; + if (p == nullptr) return -1; + if (hipPointerGetAttributes(&a, p) != hipSuccess) { + (void) hipGetLastError(); // a plain host pointer: not an error of the caller's + return -1; + } + return a.type == hipMemoryTypeDevice ? a.device : -1; +} +struct DeviceScope { + int previous = -1; + DeviceScope(int device, const char* what, const char* file, int line, size_t bytes) { + int cur = 0; + if (device < 0 || hipGetDevice(&cur) != hipSuccess || cur == device || hipSetDevice(device) != hipSuccess) return; + previous = cur; + static const bool log = [] { + const char* v = std::getenv("STRATA_HIP_XDEV_LOG"); + return v != nullptr && v[0] == '1'; + }(); + static std::atomic shown{0}; + if (log && shown.fetch_add(1) < 200) + std::fprintf(stderr, "strata hip: %s of %zu bytes on device %d while device %d was current (%s:%d)\n", what, + bytes, device, cur, file, line); + } + ~DeviceScope() { + if (previous >= 0) (void) hipSetDevice(previous); + } + DeviceScope(const DeviceScope&) = delete; + DeviceScope& operator=(const DeviceScope&) = delete; +}; +// the device side of a copy: the source of a device-to-host copy, else the destination +inline int copy_device(const void* dst, const void* src, hipMemcpyKind kind) { + if (kind == hipMemcpyDeviceToHost) return device_of(src); + const int d = device_of(dst); + return d >= 0 ? d : device_of(src); +} +inline hipError_t memcpy_on(const char* file, int line, void* dst, const void* src, size_t n, hipMemcpyKind kind) { + if (!multi_device()) return hipMemcpy(dst, src, n, kind); + const DeviceScope scope(copy_device(dst, src, kind), "memcpy", file, line, n); + return hipMemcpy(dst, src, n, kind); +} +inline hipError_t memcpy_async_on(const char* file, int line, void* dst, const void* src, size_t n, hipMemcpyKind kind, + hipStream_t stream = nullptr) { + if (stream != nullptr || !multi_device()) return hipMemcpyAsync(dst, src, n, kind, stream); + const DeviceScope scope(copy_device(dst, src, kind), "memcpy on the null stream", file, line, n); + return hipMemcpyAsync(dst, src, n, kind, stream); +} +inline hipError_t memset_on(const char* file, int line, void* dst, int value, size_t n) { + if (!multi_device()) return hipMemset(dst, value, n); + const DeviceScope scope(device_of(dst), "memset", file, line, n); + return hipMemset(dst, value, n); +} +inline hipError_t memset_async_on(const char* file, int line, void* dst, int value, size_t n, hipStream_t stream = nullptr) { + if (stream != nullptr || !multi_device()) return hipMemsetAsync(dst, value, n, stream); + const DeviceScope scope(device_of(dst), "memset on the null stream", file, line, n); + return hipMemsetAsync(dst, value, n, stream); +} +} // namespace strata::hip_compat +#define cudaMemcpy(...) strata::hip_compat::memcpy_on(__FILE__, __LINE__, __VA_ARGS__) +#define cudaMemcpyAsync(...) strata::hip_compat::memcpy_async_on(__FILE__, __LINE__, __VA_ARGS__) +#define cudaMemset(...) strata::hip_compat::memset_on(__FILE__, __LINE__, __VA_ARGS__) +#define cudaMemsetAsync(...) strata::hip_compat::memset_async_on(__FILE__, __LINE__, __VA_ARGS__) + template inline hipError_t cudaFuncSetAttribute(Kernel kernel, hipFuncAttribute attribute, int value) { return hipFuncSetAttribute(reinterpret_cast(kernel), attribute, value); diff --git a/edge0-strata/include/strata/prefill/lat_dense_w11.hpp b/edge0-strata/include/strata/prefill/lat_dense_w11.hpp new file mode 100644 index 00000000..82d67a1e --- /dev/null +++ b/edge0-strata/include/strata/prefill/lat_dense_w11.hpp @@ -0,0 +1,32 @@ +// include/strata/prefill/lat_dense_w11.hpp - the latent (edge0) D and U dense projections of the prompt path on the +// gfx11 (RDNA3 / RDNA3.5) matrix cores. HIP only: the CUDA build has latent_dense_kernel (moe_fused.cu). +// +// Both projections are Q8_0 weight rows times q8_1 activation rows. A 32-value block of each side is one exact +// s8 x s8 integer dot on v_wmma_i32_16x16x16_iu8 (two k-steps of 16), and the two block scales are applied per +// block in FP32: the same terms as iq_mmvq(Q8_0) and latent_u_grouped, in another summation order. +#pragma once + +#include + +namespace strata::prefill::dense_w11 { + +/// This build has the gfx11 kernels, the current device is one of the gfx11 targets, and `fmt` is 8 (Q8_0 +/// projections; the Q4_0 file keeps iq_mmvq / latent_u_grouped). +bool available(int fmt); + +/// D: the layer input's q8_1 rows `xq` [M, K] against the layer's G * ell Q8_0 weight rows `w`. Row t's products +/// for group g become the q8_1 row t * G + g of `outq` (ell / 32 blocks of 36 bytes), quantized as q8_1_store +/// does. `inv` (ell floats, or nullptr): the Q3_K pack's 1 / s column scales, applied before the quantization. +/// `xl_f32` (or nullptr): the same scaled products in FP32, row t * G + g of ell values, for a consumer that reads +/// the FP32 rows (STRATA_LATENT_MMQ). +void project_d(const void* w, const void* xq, int64_t M, int K, int G, int64_t ell, void* outq, float* xl_f32, + const float* inv, void* stream); + +/// U: each entry's q8_1 row of `yq` [M, K] against its label's H Q8_0 weight rows (`w`: n_groups groups of H +/// rows). `rows` [M] lists the entries bucketed by label and `goff` [n_groups + 1] the run of each label, both on +/// the device (fused::latent_dense_u_tables). `tiles`: an upper bound of the 64-entry tiles of all runs. FP32 into +/// `out` [M, H]. +void project_u(const void* w, const void* yq, const int32_t* rows, const int32_t* goff, int n_groups, int64_t tiles, + int64_t M, int K, int H, float* out, void* stream); + +} // namespace strata::prefill::dense_w11 diff --git a/edge0-strata/include/strata/prefill/lat_q3k_w11.hpp b/edge0-strata/include/strata/prefill/lat_q3k_w11.hpp new file mode 100644 index 00000000..c3bfefc0 --- /dev/null +++ b/edge0-strata/include/strata/prefill/lat_q3k_w11.hpp @@ -0,0 +1,25 @@ +// include/strata/prefill/lat_q3k_w11.hpp - the q3k pack's latent A stage (the routed experts of a prompt chunk) on +// the gfx11 matrix cores (RDNA3 / RDNA3.5, v_wmma_i32_16x16x16_iu8). It does the work of moe_fused_q3k.cu's +// latent_kernel_q3k, which needs mma.sync and is CUDA-only: gate/up (Q3_K) with SwiGLU and the q8_1 rounding of H +// in its epilogue, then down (IQ4_NL), each expert read where it already is. HIP builds only. +#pragma once + +#include "strata/prefill/moe_fused.hpp" + +#include +#include + +namespace strata::prefill::q3k_w11 { + +/// The current device runs these kernels (a gfx11 device; read once per device). +bool available(); + +/// The batch's products, as fused_q3k::experts_latent_q3k: `cnt`, `off`, `ts` and `tiles` are the device tables of +/// fused_q3k::group_latent (rows per expert, first row, first tile, and the {expert, first row} tile list). +/// `xa`: the xl rows' q8_1 blocks. `src`: entry -> xl row. `hbuf`: n rows of (iff / 64) * 80 bytes. +/// `dm`: [n][ell] FP32. `n` bounds the launch grid. +void experts(const fused::WideBatch& b, const fused::LatentGeom& g, const int32_t* cnt, const int32_t* off, + const int32_t* ts, const void* tiles, int64_t n, const void* xa, const int32_t* src, void* hbuf, float* dm, + void* stream); + +} // namespace strata::prefill::q3k_w11 diff --git a/edge0-strata/include/strata/prefill/moe_fused_q3k.hpp b/edge0-strata/include/strata/prefill/moe_fused_q3k.hpp index 926728fc..b5e84e87 100644 --- a/edge0-strata/include/strata/prefill/moe_fused_q3k.hpp +++ b/edge0-strata/include/strata/prefill/moe_fused_q3k.hpp @@ -12,6 +12,9 @@ void experts_latent_q3k(const fused::WideBatch& b, const fused::LatentGeom& g, i const void* scratch, const void* xa, const int32_t* src, void* hbuf, float* dm, void* stream); +/// This device runs experts_latent_q3k: sm_80 and newer, or a gfx11 device on HIP (lat_q3k_w11.hpp). +bool latent_fused_available(); + // The latent D projection with the pack's shared column scale (inv = 1/s, one float an ell column; nullptr: none). void latent_dense_d(const void* w, const void* xq, int64_t M, int K, int G, int64_t ell, void* outq, void* stream, int fmt, const float* inv); diff --git a/edge0-strata/include/strata/prefill/prefill.hpp b/edge0-strata/include/strata/prefill/prefill.hpp index bde04bbf..22d88e3d 100644 --- a/edge0-strata/include/strata/prefill/prefill.hpp +++ b/edge0-strata/include/strata/prefill/prefill.hpp @@ -189,7 +189,10 @@ class Prefill { bool set_peer(core::PeerExperts* peer, int64_t cap_rows, std::string& err); /// A session-level shared latent D/U pair (allocated and uploaded elsewhere, e.g. the engine): the init /// keeps no copy of its own and skips its upload; the prompt path only reads the pair. Set before the first `run`. - void set_shared_latent(uint8_t* d, uint8_t* u) { strata::core::latent_shared_d() = d; strata::core::latent_shared_u() = u; } + void set_shared_latent(int device, uint8_t* d, uint8_t* u) { + strata::core::latent_shared_d(device) = d; + strata::core::latent_shared_u(device) = u; + } /// Plan v0.3 P6: called after every chunk with the chunk's final multi-stream residual rows (device, /// T x hc*n_embd, valid until the next chunk) and the chunk's first position; the MTP draft layer builds its diff --git a/edge0-strata/src/core/native_head.cpp b/edge0-strata/src/core/native_head.cpp index bb90ed9d..0dcd7a85 100644 --- a/edge0-strata/src/core/native_head.cpp +++ b/edge0-strata/src/core/native_head.cpp @@ -122,7 +122,8 @@ NativeEmbed::~NativeEmbed() { else if (dev_) cudaFree(const_cast(dev_)); // the VRAM fallback below } -bool NativeEmbed::load(const std::vector& shards, int64_t n_embd, int64_t n_vocab, std::string& err) { +bool NativeEmbed::load(const std::vector& shards, int64_t n_embd, int64_t n_vocab, std::string& err, + bool one_device) { try { const strata::GgufModel model(shards); // --embd-gguf's one-tensor file (tools/embd_bf16_pack.py) says "strata-embd": only its tensor is checked @@ -153,7 +154,16 @@ bool NativeEmbed::load(const std::vector& shards, int64_t n_embd, i // prompt's scattered row reads go: over a PCIe Gen2 x4 link they cost more than the table's size in VRAM // is worth. Taken only with the table plus 512 MiB free; otherwise, or when the upload fails, the mapped // placement below is used as before (and the fallback is said). - if (const char* ev = std::getenv("STRATA_EMBED_VRAM"); ev != nullptr && ev[0] == '1') { + const char* ev = std::getenv("STRATA_EMBED_VRAM"); + if (ev != nullptr && ev[0] == '1' && !one_device) { + // The draft layer of a layer split gathers its rows on the last device. A table in the first device's + // VRAM cannot be read there (HIP: "an illegal memory access was encountered"), and the draft's graph + // then never finishes. Mapped host memory is readable from every device. + std::fprintf(stderr, "strata: native embedding: STRATA_EMBED_VRAM=1 not taken on a layer split: the table " + "stays in mapped host memory, which every device can read\n"); + ev = nullptr; + } + if (ev != nullptr && ev[0] == '1') { size_t free_b = 0, total_b = 0; const bool have_free = cudaMemGetInfo(&free_b, &total_b) == cudaSuccess; void* d = nullptr; diff --git a/edge0-strata/src/core/verify.cpp b/edge0-strata/src/core/verify.cpp index 1750c2eb..176caa24 100644 --- a/edge0-strata/src/core/verify.cpp +++ b/edge0-strata/src/core/verify.cpp @@ -1066,8 +1066,10 @@ bool Verifier::init(const WeightTable& wt, const ModelGeometry& g, SessionState& // every layer's D and U, resident: the graph covers all layers and touches them per layer // (the shared pair from set_shared_latent: no own copy, read-only for every verifier) if (lat_d_shared_ == nullptr) { // the session-level pair (LatentTab::load) when nobody passed it in - lat_d_shared_ = strata::core::latent_shared_d(); - lat_u_shared_ = strata::core::latent_shared_u(); + int lat_dev = 0; // this verifier's device: the pair of another device cannot be read here + cudaGetDevice(&lat_dev); + lat_d_shared_ = strata::core::latent_shared_d(lat_dev); + lat_u_shared_ = strata::core::latent_shared_u(lat_dev); } if (lat_d_shared_ == nullptr) { lat_d_ = b.take((uint64_t) g.n_layers * lat_dense_.d_bytes()); @@ -2193,11 +2195,19 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { // same path the f32 rows take). Bit-exact to the host's from_float (probe_std); the // xlseq doorbell after both copies covers it. quantize_q8_0_rows(lat_xl_, (long long) T * lat_G_, (long long) lat_ELL_, lat_act0_dev_, cs); +#if defined(STRATA_USE_HIP) // amd-hip port: HIP does not capture a D2H memcpy into the verify window's graph ("end capture: operation + // failed due to a previous error during capture"); a copy kernel writes the same bytes into the mapped host rows + strata::kernels::copy_bytes_from_mapped(m_latact0_, lat_act0_dev_, + (int64_t) ((size_t) T * (size_t) lat_G_ * ((size_t) lat_ELL_ / 32) * 34), cs); + strata::kernels::copy_bytes_from_mapped(m_latxlq_, lat_xl_, + (int64_t) ((size_t) T * (size_t) lat_G_ * (size_t) lat_ELL_ * 4), cs); +#else cudaMemcpyAsync(m_latact0_, lat_act0_dev_, (size_t) T * (size_t) lat_G_ * ((size_t) lat_ELL_ / 32) * 34, cudaMemcpyDeviceToHost, cs); cudaMemcpyAsync(m_latxlq_, lat_xl_, (size_t) T * (size_t) lat_G_ * (size_t) lat_ELL_ * 4, cudaMemcpyDeviceToHost, cs); +#endif // STRATA_USE_HIP (amd-hip port) if (trace_m_ != nullptr) gpu_stamp(trace_m_, (l * kProfPer + 35) * 2, cs); // K0-probe: D2H done doorbell_publish_value(nullptr, nullptr, nullptr, 0, 0, nullptr, nullptr, nullptr, (uint32_t*) m_xlseq_, (uint32_t) l + 1, cs); @@ -2730,6 +2740,7 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { } // #879: whether each row's logits are all finite, beside the pick (read after the window's own sync) if (nan_guard_on()) strata::kernels::logits_nonfinite_rows(head_logits_, T, (int) n_vocab_, m_nf_, cs); +#if !defined(STRATA_USE_HIP) // amd-hip port: a debug dump that does a synchronous D2H cudaMemcpy inside the window's graph capture; HIP fails the capture on it ("end capture: operation failed due to a previous error during capture") if (windows <= 12 && n_vocab_ > 0 && T > 0) { static std::vector tmp; tmp.resize((size_t) n_vocab_); @@ -2749,6 +2760,7 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { std::sqrt(s2 / (double) n_vocab_)); std::fflush(stderr); } +#endif // !STRATA_USE_HIP (amd-hip port) } stamp(g.n_layers, 1, 0); if (tap_on_ && !batch_rec_ && tap_file_ != nullptr) copy_from_mapped(m_tap_, tap_dev_, (int64_t) T * (int64_t) tap_layers_.size() * N, cs); @@ -2883,6 +2895,7 @@ bool Verifier::capture(int T, std::string& err) { err = std::string("verify: end capture: ") + cudaGetErrorString(ce); return false; } +#if !defined(STRATA_USE_HIP) // amd-hip port: this dump uses cudaGraph* names the HIP shim of this file does not map; HIP has its own listing below if (graph) { // K0: dump node order (unconditional — env was not reaching the engine) (D2H / wait_flag / H2D / kernel) size_t nn = 0; if (cudaGraphGetNodes(graph, nullptr, &nn) == cudaSuccess && nn > 0) { @@ -2909,6 +2922,7 @@ bool Verifier::capture(int T, std::string& err) { } } } +#endif // !STRATA_USE_HIP (amd-hip port) #if defined(STRATA_USE_HIP) // 143-amdk: the same listing on HIP (kernel names by hipKernelNameRefByPtr); STRATA_VERIFY_NODES=2 also lists the first nodes in graph order if (std::getenv("STRATA_VERIFY_NODES") != nullptr) { size_t nn = 0; diff --git a/edge0-strata/src/prefill/lat_dense_w11.cu b/edge0-strata/src/prefill/lat_dense_w11.cu new file mode 100644 index 00000000..ffe5121b --- /dev/null +++ b/edge0-strata/src/prefill/lat_dense_w11.cu @@ -0,0 +1,309 @@ +// src/prefill/lat_dense_w11.cu - see include/strata/prefill/lat_dense_w11.hpp. +// +// One work item is 128 weight rows against a tile of up to 64 activation rows; a persistent grid walks the items. +// A workgroup is 8 waves: 4 along the weight rows (32 each) and 2 along the tile (32 each), so a wave owns a +// 32 x 32 block of outputs as four 16 x 16 WMMA tiles. +// +// Fragments (gfx11, wave32): A lane l holds the 16 k of row l % 16, B lane l the 16 k of column l % 16, and C / D +// lane l holds D[2i + l / 16][l % 16] for i = 0..7. A is 16 weight rows and B is 16 activation rows, so a lane's +// eight results are eight weight rows of one activation row: the activation scale is the lane's own, and the eight +// weight scales come from LDS, where the item's scales are staged once. +#include "strata/prefill/lat_dense_w11.hpp" + +#include "strata/kernels/gfx_arch.hpp" +#include "strata/kernels/q8_1_finite.hpp" + +#include +#include + +#include +#include +#include +#include + +namespace strata::prefill::dense_w11 { + +namespace { + +#if defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1103__) || defined(__gfx1150__) || defined(__gfx1151__) +#define STRATA_DENSE_W11 1 +#else +#define STRATA_DENSE_W11 0 +#endif + +typedef int w11_i4 __attribute__((ext_vector_type(4))); +typedef int w11_i8 __attribute__((ext_vector_type(8))); + +constexpr int kTile = 64; // activation rows per tile +constexpr int kRows = 128; // weight rows per work item +constexpr int kThreads = 256; // 8 waves +constexpr int kMaxBlocks = 80; // 32-value blocks along K: 2560 / 32 (D reads the model width) + +void ck(cudaError_t e, const char* what) { + if (e != cudaSuccess) { + std::fprintf(stderr, "prefill latent dense (gfx11): %s: %s\n", what, cudaGetErrorString(e)); + std::exit(1); + } +} + +// signed x signed: both sides are int8 codes +__device__ __forceinline__ w11_i8 wmma_s8(w11_i4 a, w11_i4 b, w11_i8 c) { +#if STRATA_DENSE_W11 + return __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32(true, a, true, b, c, false); +#else + __builtin_trap(); + return c; +#endif +} + +// DU = true: the D projection (q8_1 out, optional FP32 out). DU = false: the U projection (FP32 out). +template +__global__ void __launch_bounds__(kThreads) +dense_w11_kernel(const uint8_t* __restrict__ w, const uint8_t* __restrict__ xq, const int32_t* __restrict__ rows, + const int32_t* __restrict__ goff, uint8_t* __restrict__ outq, float* __restrict__ outf, + const float* __restrict__ inv, int K, int ell, int G, int H, int n_groups, int64_t M, int nwork) { +#if STRATA_DENSE_W11 + const int NKB = K >> 5; // 32-value blocks along K + const int NFB = (DU ? G * ell : H) / kRows; // work items per tile + const size_t wrowb = (size_t) NKB * 34; // a Q8_0 weight row: {fp16 d; int8 qs[32]} blocks + const size_t apitch = (size_t) NKB * 36; // a q8_1 activation row: {fp16 d, fp16 sum; int8 qs[32]} + __shared__ uint16_t wsc[kRows * kMaxBlocks]; // the item's weight scales (fp16 bits), [row][block] + __shared__ int srow[kTile]; // the activation row of each tile row + const int tid = threadIdx.x, lane = tid & 31, wave = tid >> 5; + const int wm = wave & 3, wn = wave >> 2, l16 = lane & 15, hi = lane >> 4; + for (int w_ = blockIdx.x; w_ < nwork; w_ += gridDim.x) { + const int fb = w_ % NFB, t = w_ / NFB; + int base, nrows; + const uint8_t* wsec; + if constexpr (DU) { + base = kTile * t; + nrows = min(kTile, (int) (M - base)); + wsec = w; + } else { + int g = 0, tb = 0; + while (g < n_groups) { + const int nt = (goff[g + 1] - goff[g] + kTile - 1) / kTile; + if (t < tb + nt) break; + tb += nt; + ++g; + } + if (g >= n_groups) continue; // the grid bound has more tiles than the runs do + base = goff[g] + kTile * (t - tb); + nrows = min(kTile, goff[g + 1] - base); + wsec = w + (size_t) g * (size_t) H * wrowb; + } + if (nrows <= 0) continue; + const int rbase = fb * kRows; + __syncthreads(); // the previous item is done with srow / wsc + if (tid < kTile) { + const int r = base + min(tid, nrows - 1); // rows past the tile's end repeat its last one + srow[tid] = DU ? r : rows[r]; + } + { + const int row = tid >> 1; // two threads a weight row, alternate blocks each + const uint8_t* p = wsec + (size_t) (rbase + row) * wrowb; + for (int kb = tid & 1; kb < NKB; kb += 2) + wsc[row * kMaxBlocks + kb] = *reinterpret_cast(p + (size_t) kb * 34); + } + __syncthreads(); + const bool on = 32 * wn < nrows; // a wave past the tile's rows has nothing to do + float acc[2][2][8]; +#pragma unroll + for (int mt = 0; mt < 2; ++mt) +#pragma unroll + for (int nt = 0; nt < 2; ++nt) +#pragma unroll + for (int i = 0; i < 8; ++i) acc[mt][nt][i] = 0.0f; + if (on) { + const uint8_t* arow[2]; + const uint8_t* brow[2]; +#pragma unroll + for (int mt = 0; mt < 2; ++mt) arow[mt] = wsec + (size_t) (rbase + 32 * wm + 16 * mt + l16) * wrowb; +#pragma unroll + for (int nt = 0; nt < 2; ++nt) brow[nt] = xq + (size_t) srow[32 * wn + 16 * nt + l16] * apitch; + for (int kb = 0; kb < NKB; ++kb) { + w11_i4 A[2][2], B[2][2]; // [tile][k-step] + float dx[2]; +#pragma unroll + for (int mt = 0; mt < 2; ++mt) { + // a Q8_0 block is only 2-byte aligned within its row: the codes as ushort pairs + const uint16_t* q = reinterpret_cast(arow[mt] + (size_t) kb * 34 + 2); +#pragma unroll + for (int ks = 0; ks < 2; ++ks) +#pragma unroll + for (int j = 0; j < 4; ++j) + A[mt][ks][j] = (int) ((uint32_t) q[8 * ks + 2 * j] | ((uint32_t) q[8 * ks + 2 * j + 1] << 16)); + } +#pragma unroll + for (int nt = 0; nt < 2; ++nt) { + const uint32_t* q = reinterpret_cast(brow[nt] + (size_t) kb * 36); + dx[nt] = __half2float(__ushort_as_half((unsigned short) (q[0] & 0xFFFFu))); +#pragma unroll + for (int ks = 0; ks < 2; ++ks) +#pragma unroll + for (int j = 0; j < 4; ++j) B[nt][ks][j] = (int) q[1 + 4 * ks + j]; + } +#pragma unroll + for (int mt = 0; mt < 2; ++mt) { + float dw[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) + dw[i] = __half2float(__ushort_as_half(wsc[(32 * wm + 16 * mt + 2 * i + hi) * kMaxBlocks + kb])); +#pragma unroll + for (int nt = 0; nt < 2; ++nt) { + w11_i8 c = w11_i8{0, 0, 0, 0, 0, 0, 0, 0}; + c = wmma_s8(A[mt][0], B[nt][0], c); + c = wmma_s8(A[mt][1], B[nt][1], c); // the block's exact integer dots +#pragma unroll + for (int i = 0; i < 8; ++i) acc[mt][nt][i] = fmaf(dw[i] * dx[nt], (float) c[i], acc[mt][nt][i]); + } + } + } + } + if (!on) continue; + if constexpr (DU) { + // the wave's 32 weight rows are one q8_1 block of the (token, group) xl row. A lane holds 16 of a + // token's 32 values (rows 16 mt + 2 i + hi) and lane ^ 16 the other 16: the block's amax and sum are + // the pair's, the scale and the codes as q8_1_store makes them. + const int j0 = rbase + 32 * wm; // the block's first weight row + const int g_out = j0 / ell, c0 = j0 % ell; // the output group, the block's first column in it + const size_t xrowb = (size_t) (ell >> 5) * 36; // an xlq row +#pragma unroll + for (int nt = 0; nt < 2; ++nt) { + const int r = 32 * wn + 16 * nt + l16; + float v[2][8]; + float am = 0.0f, sm = 0.0f; +#pragma unroll + for (int mt = 0; mt < 2; ++mt) +#pragma unroll + for (int i = 0; i < 8; ++i) { + const int cj = c0 + 16 * mt + 2 * i + hi; + v[mt][i] = inv != nullptr ? acc[mt][nt][i] * inv[cj] : acc[mt][nt][i]; + am = fmaxf(am, fabsf(v[mt][i])); + sm += v[mt][i]; + } + am = fmaxf(am, __shfl_xor_sync(0xffffffffu, am, 16)); + sm += __shfl_xor_sync(0xffffffffu, sm, 16); + if (r >= nrows) continue; + const size_t orow = (size_t) srow[r] * (size_t) G + (size_t) g_out; + const float d = strata::kernels::q8_1_finite(am / 127.0f); + uint8_t* ob = outq + orow * xrowb + (size_t) (c0 >> 5) * 36; +#pragma unroll + for (int mt = 0; mt < 2; ++mt) +#pragma unroll + for (int i = 0; i < 8; ++i) + ob[4 + 16 * mt + 2 * i + hi] = (uint8_t) strata::kernels::q8_1_quant(v[mt][i], d, am); + if (hi == 0) *reinterpret_cast(ob) = strata::kernels::q8_1_ds(d, sm); + if (outf != nullptr) { + float* o = outf + orow * (size_t) ell + c0 + hi; +#pragma unroll + for (int mt = 0; mt < 2; ++mt) +#pragma unroll + for (int i = 0; i < 8; ++i) o[16 * mt + 2 * i] = v[mt][i]; + } + } + } else { +#pragma unroll + for (int nt = 0; nt < 2; ++nt) { + const int r = 32 * wn + 16 * nt + l16; + if (r >= nrows) continue; + float* o = outf + (size_t) srow[r] * (size_t) H + rbase + 32 * wm + hi; +#pragma unroll + for (int mt = 0; mt < 2; ++mt) +#pragma unroll + for (int i = 0; i < 8; ++i) o[16 * mt + 2 * i] = acc[mt][nt][i]; + } + } + } +#else + __builtin_trap(); +#endif +} + +struct DevInfo { + bool done = false, ok = false; + int sms = 0, occ_d = 0, occ_u = 0; +}; +std::mutex g_mu; +DevInfo g_dev[32]; + +// As moe_fused.cu's wgp_blocks: HIP counts a gfx11 WGP as one multiprocessor and the occupancy query answers one +// block where two run side by side. STRATA_PF_OCC=N sets the blocks per WGP. +int wgp_blocks(int occ) { + static const int env = [] { + const char* v = std::getenv("STRATA_PF_OCC"); + return v != nullptr ? std::atoi(v) : 0; + }(); + return env > 0 ? env : 2 * std::max(occ, 1); +} + +const DevInfo& dev_info() { + int dev = 0; + cudaGetDevice(&dev); + std::lock_guard lk(g_mu); + DevInfo& d = g_dev[dev & 31]; + if (d.done) return d; + d.done = true; +#if defined(__HIPCC__) && !defined(STRATA_HIP_GFX906) + cudaDeviceProp prop; + cudaDeviceGetAttribute(&d.sms, cudaDevAttrMultiProcessorCount, dev); + if (cudaGetDeviceProperties(&prop, dev) != cudaSuccess || !strata::kernels::gfx_arch_is_gfx11_wmma(prop.gcnArchName)) { + cudaGetLastError(); + return d; + } + if (hipOccupancyMaxActiveBlocksPerMultiprocessor(&d.occ_d, dense_w11_kernel, kThreads, 0) != hipSuccess || + hipOccupancyMaxActiveBlocksPerMultiprocessor(&d.occ_u, dense_w11_kernel, kThreads, 0) != hipSuccess) + d.occ_d = d.occ_u = 0; + cudaGetLastError(); + d.ok = d.sms > 0 && d.occ_d >= 1 && d.occ_u >= 1; +#endif + return d; +} + +void need(const DevInfo& d, int K, const char* what) { + if (!d.ok || (K >> 5) > kMaxBlocks || (K & 31) != 0) { + std::fprintf(stderr, "prefill latent dense (gfx11): %s is not available here (K = %d)\n", what, K); + std::exit(1); + } +} + +} // namespace + +bool available(int fmt) { return fmt == 8 && dev_info().ok; } + +void project_d(const void* w, const void* xq, int64_t M, int K, int G, int64_t ell, void* outq, float* xl_f32, + const float* inv, void* stream) { + const DevInfo& d = dev_info(); + need(d, K, "the D projection"); + if (M <= 0) return; + if (ell % kRows != 0) { + std::fprintf(stderr, "prefill latent dense (gfx11): the latent width %lld is not a multiple of %d\n", + (long long) ell, kRows); + std::exit(1); + } + const int64_t nwork = ((M + kTile - 1) / kTile) * (int64_t) ((size_t) G * (size_t) ell / kRows); + const unsigned grid = (unsigned) std::min(nwork, (int64_t) d.sms * wgp_blocks(d.occ_d)); + dense_w11_kernel<<>>( + (const uint8_t*) w, (const uint8_t*) xq, nullptr, nullptr, (uint8_t*) outq, xl_f32, inv, K, (int) ell, G, 0, + 0, M, (int) nwork); + ck(cudaGetLastError(), "the D projection"); +} + +void project_u(const void* w, const void* yq, const int32_t* rows, const int32_t* goff, int n_groups, int64_t tiles, + int64_t M, int K, int H, float* out, void* stream) { + const DevInfo& d = dev_info(); + need(d, K, "the U projection"); + if (M <= 0 || tiles <= 0) return; + if (H % kRows != 0) { + std::fprintf(stderr, "prefill latent dense (gfx11): the model width %d is not a multiple of %d\n", H, kRows); + std::exit(1); + } + const int64_t nwork = tiles * (int64_t) (H / kRows); + const unsigned grid = (unsigned) std::min(nwork, (int64_t) d.sms * wgp_blocks(d.occ_u)); + dense_w11_kernel<<>>( + (const uint8_t*) w, (const uint8_t*) yq, rows, goff, nullptr, out, nullptr, K, 0, 0, H, n_groups, M, + (int) nwork); + ck(cudaGetLastError(), "the U projection"); +} + +} // namespace strata::prefill::dense_w11 diff --git a/edge0-strata/src/prefill/lat_q3k_w11.cu b/edge0-strata/src/prefill/lat_q3k_w11.cu new file mode 100644 index 00000000..e592de9e --- /dev/null +++ b/edge0-strata/src/prefill/lat_q3k_w11.cu @@ -0,0 +1,403 @@ +// src/prefill/lat_q3k_w11.cu - see include/strata/prefill/lat_q3k_w11.hpp. +// +// One work item is 256 weight rows against a tile of up to 64 routed rows of one expert; a persistent grid walks +// the batch's items. Gate/up: the item is 128 features, their 128 gate rows and 128 up rows. Down: 256 output rows. +// A workgroup is 8 waves along the weight rows. A wave runs two 16-row WMMA tiles of weights against the tile's +// four 16-row activation tiles, so a weight tile is decoded once for all 64 rows. For gate/up the two weight tiles +// are the gate rows and the up rows of the same 16 features, so SwiGLU stays inside the wave. +// +// Fragments (gfx11, wave32): A lane l holds the 16 k of row l % 16, B lane l the 16 k of column l % 16, and C / D +// lane l holds D[2i + l / 16][l % 16] for i = 0..7. A k-step is 16 values: one Q3_K sub-block, and half of an +// IQ4_NL or q8_1 block. +// +// Q3_K: a weight is d * (scale - 32) * (code - 4), with one 6-bit scale per sub-block. A lane's A operand is one +// row's sub-block, so it has one scale: the lane multiplies it into its 16 codes, as (4 - code) * (scale - 32), +// which is in [-128, 124] and so an int8. The matrix core then sums a whole 32-value activation block exactly, and +// only d (one per 256 weights) and the activation block's scale are left for the FP32 step. These are the terms of +// ggml's dequantized dot product, in another summation order. +// +// The weights are decoded in registers from the GGUF blocks, where the expert already is. The item's FP32 weight +// scales go to LDS once: Q3_K's d per super-block (negated: the operand carries the other sign), IQ4_NL's d per block. +#include "strata/prefill/lat_q3k_w11.hpp" + +#include "strata/kernels/gfx_arch.hpp" +#include "strata/kernels/q8_1_finite.hpp" + +#include +#include + +#include +#include +#include +#include + +namespace strata::prefill::q3k_w11 { + +namespace { + +#if defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1103__) || defined(__gfx1150__) || defined(__gfx1151__) +#define STRATA_Q3K_W11 1 +#else +#define STRATA_Q3K_W11 0 +#endif + +typedef int w11_i4 __attribute__((ext_vector_type(4))); +typedef int w11_i8 __attribute__((ext_vector_type(8))); + +constexpr int kTile = fused::kTileRows; // routed rows per tile +constexpr int kRows = 256; // weight rows per work item +constexpr int kThreads = 256; // 8 waves +constexpr int kMaxSc = 20; // FP32 scales per weight row: Q3_K super-blocks, or IQ4_NL blocks of 640 +constexpr int kHChunk = 80; // an H row's 64-feature chunk: 64 codes, two fp16 scales, pad + +struct Tabs { + const int32_t *cnt, *off, *ts; + const int2* tiles; +}; + +void ck(cudaError_t e, const char* what) { + if (e != cudaSuccess) { + std::fprintf(stderr, "prefill latent experts (gfx11): %s: %s\n", what, cudaGetErrorString(e)); + std::exit(1); + } +} + +__device__ __forceinline__ w11_i8 wmma_s8(w11_i4 a, w11_i4 b, w11_i8 c) { +#if STRATA_Q3K_W11 + return __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32(true, a, true, b, c, false); +#else + __builtin_trap(); + return c; +#endif +} + +// 16 bytes at p as four words. A GGUF row is only 2-byte aligned (a Q3_K row is 770 B at ell = 1792): the aligned +// words around it are read and shifted, and the fifth word only when the row needs it. +__device__ __forceinline__ void ld16(const uint8_t* p, uint32_t (&w)[4]) { + const uintptr_t u = (uintptr_t) p; + const uint32_t* q = reinterpret_cast(u & ~(uintptr_t) 3); + const int sh = 8 * (int) (u & 3); + const uint32_t q4 = sh ? q[4] : 0u; + const uint32_t v[5] = {q[0], q[1], q[2], q[3], q4}; +#pragma unroll + for (int i = 0; i < 4; ++i) w[i] = (uint32_t) ((((uint64_t) v[i + 1] << 32) | v[i]) >> sh); +} +__device__ __forceinline__ float ldh(const uint8_t* p) { + return __half2float(__ushort_as_half(*reinterpret_cast(p))); +} + +// per-byte a + b (mod 256), no carry between the bytes +__device__ __forceinline__ uint32_t add8(const uint32_t a, const uint32_t b) { + return ((a & 0x7F7F7F7Fu) + (b & 0x7F7F7F7Fu)) ^ ((a ^ b) & 0x80808080u); +} + +// four IQ4_NL codes (one per byte, 0..15) to their int8 values: ggml's kvalues_iq4nl as two 8-byte tables +__device__ __forceinline__ uint32_t iq4nl_word(const uint32_t n) { + constexpr uint64_t lo = 0xF6EADDCFBFAD9881ull; // -127 -104 -83 -65 -49 -35 -22 -10 + constexpr uint64_t hi = 0x7159453526190D01ull; // 1 13 25 38 53 69 89 113 + uint32_t o = 0; +#pragma unroll + for (int c = 0; c < 4; ++c) { + const uint32_t idx = (n >> (8 * c)) & 15u; + const uint64_t t = (idx & 8u) ? hi : lo; + o |= (uint32_t) ((t >> (8 * (idx & 7u))) & 0xFFu) << (8 * c); + } + return o; +} + +// GU = true: gate/up, SwiGLU, H as q8_1 into `out`. GU = false: down, FP32 into `dm`. +template +__global__ void __launch_bounds__(kThreads) +latent_w11_kernel(const fused::WideBatch b, const Tabs tb, const fused::LatentGeom g, const uint8_t* __restrict__ act, + const int32_t* __restrict__ src, uint8_t* __restrict__ out, float* __restrict__ dm) { +#if STRATA_Q3K_W11 + const int ell = (int) g.ell, iff = (int) g.iff; + const int NKB = (GU ? ell : iff) >> 5; // 32-value blocks along K + const int NFB = GU ? iff / 128 : ell / kRows; // work items per tile + const size_t apitch = GU ? (size_t) (ell >> 5) * 36 : (size_t) (iff >> 6) * kHChunk; + const size_t hpitch = (size_t) (iff >> 6) * kHChunk; + __shared__ float wd[kRows * kMaxSc]; // the item's FP32 weight scales, [row][block] + __shared__ int srow[kTile]; // the activation row of each tile row + __shared__ float ham[GU ? 8 * kTile : 1]; // gate/up: |H| max of a wave's 16 features, [wave][tile row] + const int tid = threadIdx.x, lane = tid & 31, wv = tid >> 5, l16 = lane & 15, hi = lane >> 4; + const int t0 = tb.ts[b.e0], nwork = (tb.ts[b.e1] - t0) * NFB; + for (int w_ = blockIdx.x; w_ < nwork; w_ += gridDim.x) { + const int2 tl = tb.tiles[t0 + w_ / NFB]; + const int fb = w_ % NFB, e = tl.x, row0 = tl.y; + // off[] is a permutation for the latent grouping, so a tile's rows end at off[e] + cnt[e] + const int nrows = min(kTile, tb.cnt[e] - (row0 - tb.off[e])); + const uint8_t* blob = b.blob[e - b.e0]; + if (blob == nullptr || nrows <= 0) continue; // an expert of another launch of this layer + float s_gate = 1.0f, s_up = 1.0f, s_down = 1.0f; + if (g.tail_off) { + const float* tail = reinterpret_cast(blob + g.tail_off); + s_gate = tail[0]; + s_up = tail[1]; + s_down = tail[2]; + } + // the item's weight row `lr` (0..255). Gate/up: 0..127 the gate rows of the item's features, 128..255 their up rows + auto wrow = [&](int lr) -> const uint8_t* { + if constexpr (GU) return blob + (lr >= 128 ? g.up_off : 0) + (size_t) (128 * fb + (lr & 127)) * g.gu_row; + else return blob + g.down_off + (size_t) (kRows * fb + lr) * g.d_row; + }; + // the wave's A tile `a` (0..1), its row r16 (0..15), as an item row + auto lrow = [&](int a, int r16) -> int { + if constexpr (GU) return 128 * a + 16 * wv + r16; + else return 32 * wv + 16 * a + r16; + }; + __syncthreads(); // the previous item is done with the LDS tables + if (tid < kTile) { + const int r = row0 + min(tid, nrows - 1); // rows past the tile's end repeat its last one + srow[tid] = GU ? src[r] : r; + } + { + const uint8_t* p = wrow(tid); // a thread a weight row + if constexpr (GU) { + // block_q3_K: hmask[32], qs[64], scales[12], fp16 d. The operand is (4 - code) * (scale - 32): -d here + for (int sb = 0; sb < (NKB >> 3); ++sb) wd[tid * kMaxSc + sb] = -ldh(p + (size_t) 110 * sb + 108); + } else { + for (int kb = 0; kb < NKB; ++kb) wd[tid * kMaxSc + kb] = ldh(p + (size_t) 18 * kb); + } + } + __syncthreads(); + float acc[2][4][8]; +#pragma unroll + for (int a = 0; a < 2; ++a) +#pragma unroll + for (int nt = 0; nt < 4; ++nt) +#pragma unroll + for (int i = 0; i < 8; ++i) acc[a][nt][i] = 0.0f; + const uint8_t* arow[2]; + const uint8_t* brow[4]; +#pragma unroll + for (int a = 0; a < 2; ++a) arow[a] = wrow(lrow(a, l16)); +#pragma unroll + for (int nt = 0; nt < 4; ++nt) brow[nt] = act + (size_t) srow[16 * nt + l16] * apitch; + float dw[2][8]; // the FP32 scales of the lane's weight rows + uint32_t sw[2][3]; // Q3_K: the super-block's 12 scale bytes + for (int kb = 0; kb < NKB; ++kb) { + if constexpr (GU) { + if ((kb & 7) == 0) { // a new 256-weight super-block +#pragma unroll + for (int a = 0; a < 2; ++a) { + uint32_t t[4]; + ld16(arow[a] + (size_t) 110 * (kb >> 3) + 96, t); + sw[a][0] = t[0]; + sw[a][1] = t[1]; + sw[a][2] = t[2]; +#pragma unroll + for (int i = 0; i < 8; ++i) dw[a][i] = wd[lrow(a, 2 * i + hi) * kMaxSc + (kb >> 3)]; + } + } + } else { +#pragma unroll + for (int a = 0; a < 2; ++a) +#pragma unroll + for (int i = 0; i < 8; ++i) dw[a][i] = wd[lrow(a, 2 * i + hi) * kMaxSc + kb]; + } + w11_i4 A[2][2]; // [tile][k-step] +#pragma unroll + for (int a = 0; a < 2; ++a) { + if constexpr (GU) { + // ggml's dequantize_row_q3_K: sub-block sb of the super-block reads 16 bytes of qs at + // 32 * half + 16 * gg with the 2-bit field j, and bit j + 4 * half of 16 bytes of hmask + const uint8_t* q = arow[a] + (size_t) 110 * (kb >> 3); +#pragma unroll + for (int ks = 0; ks < 2; ++ks) { + const int sb = 2 * (kb & 7) + ks, half = sb >> 3, j = (sb >> 1) & 3, gg = sb & 1, hb = j + 4 * half; + const uint32_t byte = (sw[a][(sb & 7) >> 2] >> (8 * (sb & 3))) & 0xFFu; + const uint32_t nib = (byte >> (4 * (sb >> 3))) & 0xFu; + const uint32_t hi2 = (sw[a][2] >> (8 * (sb & 3) + 2 * (sb >> 2))) & 3u; + const uint32_t S = (uint32_t) ((int) (nib | (hi2 << 4)) - 32) & 0xFFu; // scale - 32, a byte + // (4 - code) * S = c * S - 3 S with c = 7 - code: the three bit planes of c select S, 2 S, 4 S + const uint32_t r0 = S * 0x01010101u, r1 = ((2u * S) & 0xFFu) * 0x01010101u, + r2 = ((4u * S) & 0xFFu) * 0x01010101u, rk = ((0u - 3u * S) & 0xFFu) * 0x01010101u; + uint32_t qv[4], hv[4]; + ld16(q + 32 + 32 * half + 16 * gg, qv); + ld16(q + 16 * gg, hv); +#pragma unroll + for (int m = 0; m < 4; ++m) { + const uint32_t c = (((qv[m] >> (2 * j)) & 0x03030303u) | (((hv[m] >> hb) & 0x01010101u) << 2)) ^ + 0x07070707u; + const uint32_t p0 = ((c & 0x01010101u) * 0xFFu) & r0; + const uint32_t p1 = (((c >> 1) & 0x01010101u) * 0xFFu) & r1; + const uint32_t p2 = (((c >> 2) & 0x01010101u) * 0xFFu) & r2; + A[a][ks][m] = (int) add8(add8(p0, p1), add8(p2, rk)); + } + } + } else { + // an IQ4_NL block: fp16 d, qs[16]. Value j is the low nibble of qs[j], value 16 + j the high. + // The block is 2-byte aligned: qs is either word aligned, or it starts 2 bytes into a word + const uint8_t* q = arow[a] + (size_t) 18 * kb; + uint32_t n[4]; + if ((uintptr_t) q & 2) { + const uint32_t* c = reinterpret_cast(q + 2); +#pragma unroll + for (int m = 0; m < 4; ++m) n[m] = c[m]; + } else { + const uint32_t* c = reinterpret_cast(q); + const uint32_t v[5] = {c[0], c[1], c[2], c[3], *reinterpret_cast(q + 16)}; +#pragma unroll + for (int m = 0; m < 4; ++m) n[m] = (v[m] >> 16) | (v[m + 1] << 16); + } +#pragma unroll + for (int m = 0; m < 4; ++m) { + A[a][0][m] = (int) iq4nl_word(n[m] & 0x0F0F0F0Fu); + A[a][1][m] = (int) iq4nl_word((n[m] >> 4) & 0x0F0F0F0Fu); + } + } + } +#pragma unroll + for (int nt = 0; nt < 4; ++nt) { + if (16 * nt >= nrows) continue; // the tile has no rows here + w11_i4 B[2]; + float dx; + if constexpr (GU) { // a q8_1 block: {fp16 d, fp16 sum; int8 qs[32]} + const uint32_t* q = reinterpret_cast(brow[nt] + (size_t) 36 * kb); + dx = __half2float(__ushort_as_half((unsigned short) (q[0] & 0xFFFFu))); +#pragma unroll + for (int ks = 0; ks < 2; ++ks) +#pragma unroll + for (int j = 0; j < 4; ++j) B[ks][j] = (int) q[1 + 4 * ks + j]; + } else { // an H chunk: 64 codes, then the two blocks' fp16 d + const uint8_t* q = brow[nt] + (size_t) kHChunk * (kb >> 1); + const uint32_t* c = reinterpret_cast(q + 32 * (kb & 1)); + dx = __half2float(__ushort_as_half(*reinterpret_cast(q + 64 + 2 * (kb & 1)))); +#pragma unroll + for (int ks = 0; ks < 2; ++ks) +#pragma unroll + for (int j = 0; j < 4; ++j) B[ks][j] = (int) c[4 * ks + j]; + } +#pragma unroll + for (int a = 0; a < 2; ++a) { + w11_i8 c = w11_i8{0, 0, 0, 0, 0, 0, 0, 0}; + c = wmma_s8(A[a][0], B[0], c); + c = wmma_s8(A[a][1], B[1], c); // the block's exact integer dots +#pragma unroll + for (int i = 0; i < 8; ++i) acc[a][nt][i] = fmaf(dw[a][i] * dx, (float) c[i], acc[a][nt][i]); + } + } + } + if constexpr (GU) { + // a q8_1 block of H is 32 features: this wave's 16 and those of wave ^ 1. A lane holds 8 of a tile row's + // features (2 i + hi) and lane ^ 16 the other 8. The scale and the codes are as q8_1_store makes them. + float h[4][8]; +#pragma unroll + for (int nt = 0; nt < 4; ++nt) { + float am = 0.0f; +#pragma unroll + for (int i = 0; i < 8; ++i) { + const float gt = acc[0][nt][i] * s_gate, up = acc[1][nt][i] * s_up; + h[nt][i] = gt / (1.0f + expf(-gt)) * up; + am = fmaxf(am, fabsf(h[nt][i])); + } + am = fmaxf(am, __shfl_xor_sync(0xffffffffu, am, 16)); + if (hi == 0) ham[wv * kTile + 16 * nt + l16] = am; + } + __syncthreads(); + const int f0 = 16 * wv; // the wave's first feature in the item + const int hh = (f0 >> 5) & 1; // its block in the 64-feature chunk +#pragma unroll + for (int nt = 0; nt < 4; ++nt) { + const int r = 16 * nt + l16; + if (r >= nrows) continue; + const float am = fmaxf(ham[wv * kTile + r], ham[(wv ^ 1) * kTile + r]); + const float d = strata::kernels::q8_1_finite(am / 127.0f); + uint8_t* o = out + (size_t) (row0 + r) * hpitch + (size_t) (2 * fb + (f0 >> 6)) * kHChunk; +#pragma unroll + for (int i = 0; i < 8; ++i) + o[(f0 & 63) + 2 * i + hi] = (uint8_t) strata::kernels::q8_1_quant(h[nt][i], d, am); + if (hi == 0 && (wv & 1) == 0) *reinterpret_cast<__half*>(o + 64 + 2 * hh) = __float2half(d); + } + } else { +#pragma unroll + for (int nt = 0; nt < 4; ++nt) { + const int r = 16 * nt + l16; + if (r >= nrows) continue; + float* o = dm + (size_t) (row0 + r) * (size_t) ell + kRows * fb + 32 * wv + hi; +#pragma unroll + for (int a = 0; a < 2; ++a) +#pragma unroll + for (int i = 0; i < 8; ++i) o[16 * a + 2 * i] = acc[a][nt][i] * s_down; + } + } + } +#else + __builtin_trap(); +#endif +} + +struct DevInfo { + bool done = false, ok = false; + int sms = 0, occ_gu = 0, occ_d = 0; +}; +std::mutex g_mu; +DevInfo g_dev[32]; + +// As moe_fused.cu's wgp_blocks: HIP counts a gfx11 WGP as one multiprocessor and the occupancy query answers one +// block where two run side by side. STRATA_PF_OCC=N sets the blocks per WGP. +int wgp_blocks(int occ) { + static const int env = [] { + const char* v = std::getenv("STRATA_PF_OCC"); + return v != nullptr ? std::atoi(v) : 0; + }(); + return env > 0 ? env : 2 * std::max(occ, 1); +} + +const DevInfo& dev_info() { + int dev = 0; + cudaGetDevice(&dev); + std::lock_guard lk(g_mu); + DevInfo& d = g_dev[dev & 31]; + if (d.done) return d; + d.done = true; +#if defined(__HIPCC__) && !defined(STRATA_HIP_GFX906) + cudaDeviceProp prop; + cudaDeviceGetAttribute(&d.sms, cudaDevAttrMultiProcessorCount, dev); + if (cudaGetDeviceProperties(&prop, dev) != cudaSuccess || !strata::kernels::gfx_arch_is_gfx11_wmma(prop.gcnArchName)) { + cudaGetLastError(); + return d; + } + if (hipOccupancyMaxActiveBlocksPerMultiprocessor(&d.occ_gu, latent_w11_kernel, kThreads, 0) != hipSuccess || + hipOccupancyMaxActiveBlocksPerMultiprocessor(&d.occ_d, latent_w11_kernel, kThreads, 0) != hipSuccess) + d.occ_gu = d.occ_d = 0; + cudaGetLastError(); + d.ok = d.sms > 0 && d.occ_gu >= 1 && d.occ_d >= 1; +#endif + return d; +} + +} // namespace + +bool available() { return dev_info().ok; } + +void experts(const fused::WideBatch& b, const fused::LatentGeom& g, const int32_t* cnt, const int32_t* off, + const int32_t* ts, const void* tiles, int64_t n, const void* xa, const int32_t* src, void* hbuf, float* dm, + void* stream) { + if (b.e1 <= b.e0 || n <= 0) return; + const DevInfo& d = dev_info(); + // the geometry these kernels take: Q3_K gate/up rows of whole 256-value super-blocks, IQ4_NL down rows + const bool geom = d.ok && g.ell > 0 && g.ell % 256 == 0 && (g.ell >> 8) <= kMaxSc && g.iff > 0 && g.iff % 128 == 0 && + (g.iff >> 5) <= kMaxSc && g.gu_row == (size_t) (g.ell / 256) * 110 && + g.d_row == (size_t) (g.iff / 32) * 18; + if (!geom) { + std::fprintf(stderr, + "prefill latent experts (gfx11): not available here (ell %lld, iff %lld, rows of %zu and %zu bytes)\n", + (long long) g.ell, (long long) g.iff, g.gu_row, g.d_row); + std::exit(1); + } + const cudaStream_t s = (cudaStream_t) stream; + const Tabs tb{cnt, off, ts, (const int2*) tiles}; + // the most tiles the batch can have: every row in it, plus a partial tile per expert + const int64_t tiles_max = (n + kTile - 1) / kTile + (b.e1 - b.e0); + const unsigned g_gu = + (unsigned) std::min(tiles_max * (g.iff / 128), (int64_t) d.sms * wgp_blocks(d.occ_gu)); + const unsigned g_d = + (unsigned) std::min(tiles_max * (g.ell / kRows), (int64_t) d.sms * wgp_blocks(d.occ_d)); + latent_w11_kernel<<>>(b, tb, g, (const uint8_t*) xa, src, (uint8_t*) hbuf, nullptr); + ck(cudaGetLastError(), "gate/up"); + latent_w11_kernel<<>>(b, tb, g, (const uint8_t*) hbuf, src, nullptr, dm); + ck(cudaGetLastError(), "down"); +} + +} // namespace strata::prefill::q3k_w11 diff --git a/edge0-strata/src/prefill/moe_fused.cu b/edge0-strata/src/prefill/moe_fused.cu index da5d5906..8b38d8d8 100644 --- a/edge0-strata/src/prefill/moe_fused.cu +++ b/edge0-strata/src/prefill/moe_fused.cu @@ -1407,11 +1407,13 @@ void experts_latent(const WideBatch& b, const LatentGeom& g, int n_expert, int64 const int64_t tiles = (n + kTileRows - 1) / kTileRows + (b.e1 - b.e0); const unsigned g_gu = (unsigned) std::min(tiles * 5, (int64_t) d.sms * d.occ_gu); const unsigned g_d = (unsigned) std::min(tiles * 7, (int64_t) d.sms * d.occ_d); +#if !defined(__HIPCC__) // amd-hip port: the latent A kernels are CUDA-only (lat_dev().ok is false on HIP, so this is not reached) latent_kernel<<>>(b, tb, g, (const uint8_t*) xa, src, (uint8_t*) hbuf, nullptr); ck(cudaGetLastError(), "experts_latent gu"); latent_kernel<<>>(b, tb, g, (const uint8_t*) hbuf, src, nullptr, dm); ck(cudaGetLastError(), "experts_latent down"); +#endif // !__HIPCC__ (amd-hip port) ck(cudaGetLastError(), "experts_latent"); } diff --git a/edge0-strata/src/prefill/moe_fused_q3k.cu b/edge0-strata/src/prefill/moe_fused_q3k.cu index 3bcf3356..021fd982 100644 --- a/edge0-strata/src/prefill/moe_fused_q3k.cu +++ b/edge0-strata/src/prefill/moe_fused_q3k.cu @@ -34,6 +34,9 @@ #define GGML_COMMON_DECL_CUDA #define GGML_COMMON_IMPL_CUDA #include "ggml-common.h" // kvalues_iq4nl, block_q8_1 +#if defined(__HIPCC__) +#include "strata/prefill/lat_q3k_w11.hpp" // the latent A stage on the gfx11 matrix cores +#endif #include #include @@ -1481,7 +1484,13 @@ void experts(const Batch& b, int n_expert, int64_t n, const void* scratch, const ck(cudaGetLastError(), "experts"); } -bool latent_fused_available() { return lat_dev().ok; } +bool latent_fused_available() { +#if defined(__HIPCC__) + return q3k_w11::available(); // latent_kernel_q3k needs mma.sync: gfx11 has its own kernels +#else + return lat_dev().ok; +#endif +} size_t latent_group_bytes(int64_t n, int n_expert) { return group_bytes(n, n_expert); } @@ -1521,6 +1530,13 @@ void group_latent(const int32_t* cnt, const int32_t* off, int n_expert, void* ho void experts_latent_q3k(const WideBatch& b, const LatentGeom& g, int n_expert, int64_t n, const void* scratch, const void* xa, const int32_t* src, void* hbuf, float* dm, void* stream) { if (b.e1 <= b.e0 || n <= 0) return; +#if defined(__HIPCC__) + { + const Tables tb = tables(const_cast(scratch), n_expert); + q3k_w11::experts(b, g, tb.cnt, tb.off, tb.ts, tb.tiles, n, xa, src, hbuf, dm, stream); + return; + } +#endif const LatentDevInfo& d = lat_dev(); if (!d.ok) { std::fprintf(stderr, "prefill fused experts (latent): the kernels are not available on this device\n"); @@ -1534,11 +1550,13 @@ void experts_latent_q3k(const WideBatch& b, const LatentGeom& g, int n_expert, i const int64_t nfb_d = (int64_t) std::max(1, (int64_t) (g.ell >> 8)); const unsigned g_gu = (unsigned) std::min(tiles * nfb_gu, (int64_t) d.sms * d.occ_gu); const unsigned g_d = (unsigned) std::min(tiles * nfb_d, (int64_t) d.sms * d.occ_d); +#if !defined(__HIPCC__) // amd-hip port: the latent A kernels are CUDA-only (lat_dev().ok is false on HIP, so this is not reached) latent_kernel_q3k<<>>(b, tb, g, (const uint8_t*) xa, src, (uint8_t*) hbuf, nullptr); ck(cudaGetLastError(), "experts_latent_q3k gu"); latent_kernel_q3k<<>>(b, tb, g, (const uint8_t*) hbuf, src, nullptr, dm); ck(cudaGetLastError(), "experts_latent_q3k down"); +#endif // !__HIPCC__ (amd-hip port) ck(cudaGetLastError(), "experts_latent_q3k"); } diff --git a/edge0-strata/src/prefill/prefill.cpp b/edge0-strata/src/prefill/prefill.cpp index bf1db352..a20f2cb8 100644 --- a/edge0-strata/src/prefill/prefill.cpp +++ b/edge0-strata/src/prefill/prefill.cpp @@ -52,6 +52,9 @@ #include "strata/prefill/moe_mmq.hpp" #include "strata/prefill/lat_group_tables.hpp" #include "strata/prefill/lat_q3k_fused.hpp" // M4: STRATA_LAT_FUSED_Q3K +#if defined(STRATA_USE_HIP) +#include "strata/prefill/lat_dense_w11.hpp" // the latent D / U projections on the gfx11 matrix cores +#endif #include "strata/prefill/moe_fused_q3k.hpp" // the q3k pack's variant of v5's latent kernel #include "strata/core/peer_experts.hpp" #include "strata/core/expert_nibble.hpp" @@ -1729,8 +1732,10 @@ bool Prefill::init(const core::WeightTable& wt, const core::ModelGeometry& g, co m.owned.push_back(m.lat_labels_dev); } if (m.latent && m.lat_d_dev == nullptr) { // a session-shared pair when present (latent_share.hpp), else owned below - m.lat_d_dev = strata::core::latent_shared_d(); - m.lat_u_dev = strata::core::latent_shared_u(); + int lat_dev = 0; // this prompt path's device: the pair of another device cannot be read here + cudaGetDevice(&lat_dev); + m.lat_d_dev = strata::core::latent_shared_d(lat_dev); + m.lat_u_dev = strata::core::latent_shared_u(lat_dev); } if (m.latent && m.lat_d_dev == nullptr) { // no shared pair: one owned copy, uploaded once per session const uint64_t db = (uint64_t) g.n_layers * m.lat_d.d_bytes(); @@ -4988,7 +4993,12 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri const char* v = std::getenv("STRATA_LATENT_FUSED_DU"); return v == nullptr || v[0] != '0'; }(); +#if defined(STRATA_USE_HIP) // gfx11: the D / U projections have their own WMMA kernels (lat_dense_w11.cu); the + // fused A stage stays CUDA-only, so latent_fused_available() is false here + const bool lat_du = lat_du_env && dense_w11::available(m.lat_fmt); +#else const bool lat_du = lat_du_env && fused::latent_fused_available(); +#endif // D: each token's 8 group products [T, G*ELL] -> q8_1 rows (xlq) pt.mark(kPfGather, m.cs); strata::kernels::quantize_q8_1_rows(m.mixed, T, (int) N, m.lat_xq, m.cs); @@ -4997,12 +5007,23 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri // DU kernel that folds 1/s into the xl columns; v5's pack has no scales and keeps // the original kernel (its signature and its bits are untouched) const float* lat_inv = strata::kernels::lat_gu_colscale_dev((int) l, (int) ELL); +#if defined(STRATA_USE_HIP) + // STRATA_LATENT_MMQ's A stage quantizes from the FP32 xl rows, so they are written + // too when it is asked for; the native A stage reads xlq only + static const bool lat_xl_f32 = [] { + const char* v = std::getenv("STRATA_LATENT_MMQ"); + return v != nullptr && v[0] == '1'; + }(); + dense_w11::project_d(m.lat_d_dev + (size_t) l * m.lat_d.d_bytes(), m.lat_xq, T, (int) N, + (int) G, ELL, m.xlq, lat_xl_f32 ? m.xl : nullptr, lat_inv, m.cs); +#else if (lat_inv) fused_q3k::latent_dense_d(m.lat_d_dev + (size_t) l * m.lat_d.d_bytes(), m.lat_xq, T, (int) N, (int) G, ELL, m.xlq, m.cs, m.lat_fmt, lat_inv); else fused::latent_dense_d(m.lat_d_dev + (size_t) l * m.lat_d.d_bytes(), m.lat_xq, T, (int) N, (int) G, ELL, m.xlq, m.cs, m.lat_fmt); +#endif } else { strata::kernels::iq_mmvq(strata::kernels::lat_dense_fmt_to_type(m.lat_fmt), m.lat_d_dev + (size_t) l * m.lat_d.d_bytes(), m.lat_xq, m.xl, (int) N, @@ -5108,10 +5129,18 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri const char* v = std::getenv("STRATA_LATENT_FUSED_LOG"); return v != nullptr && v[0] == '1'; }(); +#if defined(STRATA_USE_HIP) + // HIP: the q3k pack's A stage has gfx11 kernels of its own (lat_q3k_w11); the NVFP4 + // pack's fused kernels are CUDA-only + const bool lat_fused = + lat_fused_env && !nib_l && lat_fused_q3k_on() && f.gu_type == 11 && f.d_type == 20 && + m.lat_I == 640 && fused_q3k::latent_fused_available(); +#else const bool lat_fused = lat_fused_env && !nib_l && fused::latent_fused_available() && ((f.gu_type == 40 && f.d_type == 40) || (lat_fused_q3k_on() && f.gu_type == 11 && f.d_type == 20 && m.lat_I == 640)); +#endif fused::LatentGeom lat_lg{}; size_t lat_h = 0; if (lat_fused) { @@ -5151,9 +5180,13 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri m.grp_d != nullptr && m.bounds_dev != nullptr && !nib_l; // M4: STRATA_LAT_FUSED_Q3K=1 - the q3k pack only (Q3_K gate/up, IQ4_NL down, the 1280 / // 640 latent geometry), and it takes the launches in place of both other paths. +#if defined(STRATA_USE_HIP) + const bool lat_fused_q3k = false; // lat_q3k_fused's kernels are CUDA-only (stubs that exit) +#else const bool lat_fused_q3k = lat_fused_q3k_on() && f.gu_type == 11 /* GGML_TYPE_Q3_K */ && f.d_type == 20 /* GGML_TYPE_IQ4_NL */ && ELL == 1280 && m.lat_I == 640 && m.lat_hbuf != nullptr && m.lat_gcd != nullptr && !nib_l; +#endif if (lat_fused_q3k && !nib_l) { static bool noted = false; if (!noted) { @@ -5338,6 +5371,7 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri for (size_t q = 0; q <= nb2 && q < 17; ++q) std::fprintf(stderr, "latmmq-bh[%zu]=%d\n", q, bh[q]); } +#if !defined(STRATA_USE_HIP) // amd-hip port: trace-only block (STRATA_LATENT_MMQ_TRACE), uses cudaPointerGetAttributes which the HIP shim does not map { // latmmq-dev: which device each buffer lives on, and the current one static bool devsaid = false; if (!devsaid && std::getenv("STRATA_LATENT_MMQ_TRACE") != nullptr) { @@ -5366,6 +5400,7 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri devof(m.lat_gdst)); } } +#endif // !STRATA_USE_HIP (amd-hip port) auto run_p = [&](const mmq::Product& pr) { static const bool alt = std::getenv("STRATA_LATENT_MMQ_ALTCTX") != nullptr; static std::unique_ptr altctx; @@ -5906,11 +5941,18 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri (int32_t*) m.lat_scratch, (int32_t*) m.lat_scratch + 12, (int32_t*) m.lat_scratch + 32, m.cs); // the grid bound without the counts: every split fits in ceil(n / 64) + G tiles +#if defined(STRATA_USE_HIP) + dense_w11::project_u(m.lat_u_dev + (size_t) l * m.lat_d.u_bytes(), m.ylq, + (const int32_t*) m.lat_scratch + 32, (const int32_t*) m.lat_scratch, + (int) G, ((int64_t) (T * K) + 63) / 64 + G, (int64_t) (T * K), + (int) ELL, (int) N, m.Dm, m.cs); +#else const int64_t unw = (((int64_t) (T * K) + 63) / 64 + G) * (N >> 8); fused::latent_dense_u(m.lat_u_dev + (size_t) l * m.lat_d.u_bytes(), m.ylq, (const int32_t*) m.lat_scratch + 32, (const int32_t*) m.lat_scratch, (int) G, unw, (int64_t) (T * K), (int) ELL, (int) N, m.Dm, m.cs, m.lat_fmt); +#endif } else { strata::kernels::latent_u_grouped(m.lat_u_dev + (size_t) l * m.lat_d.u_bytes(), m.ylq, m.ids_p, m.lat_labels_dev + (size_t) l * m.lat_E, m.lat_E, (int64_t) (T * K), diff --git a/edge0-strata/src/program/generate.cpp b/edge0-strata/src/program/generate.cpp index 4a2b3fbe..0e4e5f54 100644 --- a/edge0-strata/src/program/generate.cpp +++ b/edge0-strata/src/program/generate.cpp @@ -232,8 +232,17 @@ struct LatentTab { int64_t lat_H = 0, lat_ELL = 0, lat_I = 0; int lat_G = 8, lat_E = 0; bool latent = false; - bool load(const strata::core::ModelGeometry& g, std::string& err) { - if (d != nullptr) { strata::core::latent_shared_d() = d; strata::core::latent_shared_u() = u; return true; } + bool opened = false; + // The pair on the CURRENT device, for layers [lb, le) (le < 0: every layer). A layer split's later stage calls + // this on its own device: device memory cannot be read from another device, so each device has its own pair. + // A stage's pair holds only its layers; the pointers are offset so that layer l is at d + l * d_bytes. + bool load(const strata::core::ModelGeometry& g, std::string& err, int64_t lb = 0, int64_t le = -1) { + int dev = 0; + cudaGetDevice(&dev); + if (strata::core::latent_shared_d(dev) != nullptr) return true; + if (le < 0 || le > g.n_layers) le = g.n_layers; + if (lb < 0 || lb >= le) lb = 0; + if (opened) return upload(g, dev, lb, le, err); const auto& lay0 = strata::kernels::cpu::expert_layout(); if (!(lay0.native && !lay0.fmt.empty() && lay0.fmt[0].n_embd != (int64_t) g.n_embd)) return true; latent = true; @@ -257,22 +266,32 @@ struct LatentTab { std::fclose(lf); if (got != labels.size()) { err = "latent dense: " + lp + " is too small"; return false; } } - const uint64_t db = (uint64_t) g.n_layers * dense.d_bytes(), ub = (uint64_t) g.n_layers * dense.u_bytes(); - if (cudaMalloc((void**) &d, db) != cudaSuccess || cudaMalloc((void**) &u, ub) != cudaSuccess) { + opened = true; + return upload(g, dev, lb, le, err); + } + bool upload(const strata::core::ModelGeometry& g, int dev, int64_t lb, int64_t le, std::string& err) { + (void) g; + const uint64_t db = (uint64_t) (le - lb) * dense.d_bytes(), ub = (uint64_t) (le - lb) * dense.u_bytes(); + uint8_t *dd = nullptr, *du = nullptr; + if (cudaMalloc((void**) &dd, db) != cudaSuccess || cudaMalloc((void**) &du, ub) != cudaSuccess) { const cudaError_t ce = cudaGetLastError(); - if (d) { cudaFree(d); d = nullptr; } - if (u) { cudaFree(u); u = nullptr; } + if (dd) cudaFree(dd); + if (du) cudaFree(du); err = std::string("the latent D/U device buffers could not be had (") + cudaGetErrorString(ce) + ")"; return false; } - for (int64_t l = 0; l < g.n_layers; ++l) - if (cudaMemcpy(d + (size_t) l * dense.d_bytes(), dense.d(l), dense.d_bytes(), cudaMemcpyHostToDevice) != cudaSuccess || - cudaMemcpy(u + (size_t) l * dense.u_bytes(), dense.u(l), dense.u_bytes(), cudaMemcpyHostToDevice) != cudaSuccess) { + for (int64_t l = lb; l < le; ++l) + if (cudaMemcpy(dd + (size_t) (l - lb) * dense.d_bytes(), dense.d(l), dense.d_bytes(), cudaMemcpyHostToDevice) != cudaSuccess || + cudaMemcpy(du + (size_t) (l - lb) * dense.u_bytes(), dense.u(l), dense.u_bytes(), cudaMemcpyHostToDevice) != cudaSuccess) { err = "the latent D/U upload failed at layer " + std::to_string(l); return false; } - strata::core::latent_shared_d() = d; - strata::core::latent_shared_u() = u; + // layer l is read at base + l * bytes: a stage's pair starts at its first layer + strata::core::latent_shared_d(dev) = dd - (size_t) lb * dense.d_bytes(); + strata::core::latent_shared_u(dev) = du - (size_t) lb * dense.u_bytes(); + if (d == nullptr) { d = strata::core::latent_shared_d(dev); u = strata::core::latent_shared_u(dev); } + std::fprintf(stderr, "strata generate: latent D/U on CUDA%d: layers %lld-%lld, %.2f GiB\n", dev, (long long) lb, + (long long) le - 1, (double) (db + ub) / 1073741824.0); return true; } } lat_tab; @@ -4486,7 +4505,7 @@ int main(int argc, char** argv) { } const auto embed_t0 = std::chrono::steady_clock::now(); if (!native_embed.load(o.embd_gguf.empty() ? o.native_shards : std::vector{o.embd_gguf}, g.n_embd, - g.n_vocab, err)) { + g.n_vocab, err, !multi_gpu)) { std::fprintf(stderr, "strata generate: %s\n", err.c_str()); return 1; } @@ -6046,6 +6065,8 @@ int main(int argc, char** argv) { for (size_t i = 0; i < stages.size(); ++i) { GpuStage& st = *stages[i]; const strata::core::OnDevice on(st.dev); + // latent MoE: this stage's own D/U pair, for the layers it runs, before its expert cache is sized + { std::string le; if (!lat_tab.load(g, le, st.lb, st.le)) { std::fprintf(stderr, "strata generate: layer split, CUDA%d: %s\n", st.dev, le.c_str()); return 1; } } void* sbuf_s = nullptr; if (cudaMalloc(&sbuf_s, strata::core::session_bytes(g, o.max_context, K, st.lb, st.le)) != cudaSuccess || strata::core::session_init(g, o.max_context, K, sbuf_s, st.ss, st.lb, st.le) == 0) { @@ -6924,7 +6945,8 @@ int main(int argc, char** argv) { }; if (o.expert_cache < 0) { // latent MoE: the session D/U pair is loaded before the free-VRAM read, so auto leaves room for it - { std::string le; if (!lat_tab.load(g, le)) { std::fprintf(stderr, "strata generate: %s\n", le.c_str()); return 1; } } + // (a layer split: only the first stage's layers, the later stages hold their own) + { std::string le; if (!lat_tab.load(g, le, 0, multi_gpu && !split_at.empty() ? split_at[0] : -1)) { std::fprintf(stderr, "strata generate: %s\n", le.c_str()); return 1; } } size_t free_b = 0, total_b = 0; free_b = strata::core::device_free_bytes(); (void) total_b; if (free_b == 0 && o.prefill_head_alias) { @@ -10250,7 +10272,7 @@ int main(int argc, char** argv) { for (int st = 0; st + 1 < n_stages; ++st) stage_ver_p(st, 1).set_next(&stage_ver_p(st + 1, 1), &split_drive_b); } ver.set_remote_expert_opt(remote_opt.get()); - { std::string le; if (!lat_tab.load(g, le)) { std::fprintf(stderr, "strata serve: %s\n", le.c_str()); return 1; } } + { std::string le; if (!lat_tab.load(g, le, 0, multi_gpu && !split_at.empty() ? split_at[0] : -1)) { std::fprintf(stderr, "strata serve: %s\n", le.c_str()); return 1; } } if (pcie_balance() && n_stages == 1) ver.set_pcie_balance(&drive.d.pcie_model); if (batch_mtp) ver.set_batch_graph_limit(64); // --batch-mtp only: slot rotation makes many layouts (LRU); 0.1.39 keeps all if (!ver.init(wt, g, ss, vh, native_head.loaded() || prefill_head_alias_bound ? &native_head : nullptr,