diff --git a/.agents/issue-index.md b/.agents/issue-index.md index 5fc4b1cd8..6740fe47a 100644 --- a/.agents/issue-index.md +++ b/.agents/issue-index.md @@ -372,3 +372,4 @@ rather than merged. `scripts/check-agent-record.py` gates both. | [#1190](https://github.com/mudler/vllm.cpp/issues/1190) | `ENV-AGNOSTIC-CAMPAIGN` | One operator's hosts, share paths and addresses are written into 227 tracked files where a `.env` placeholder belongs, so a second developer who follows the protocol documents is told to reach a box on one home network. Re-derived at `fd64c76ee` with `git grep -cIE 'dgx\.casa\|nas_share\|192\.168\.\|thor:gpu0'`: `dgx.casa` 203 files, `nas_share` 32, `192.168.` 24, `thor:gpu0` 11. The mechanism was already there and almost unused: `.env.example` is tracked, `.env` and `.agents/developer-preferences.md` are ignored, `scripts/agent-onboard.py --env-set` already records one answered value and refuses an undeclared key, and `${VLLM_SOURCE}` and `${GPU_LOCK}` already resolve in 59 and 50 files while `${DEVICE_ARCH}`, `${DEVICE_TOOLKIT_ROOT}` and `${DEVICE_COMPILER}` resolve in none. This SCOPING row lands the rule and the mechanism, not the sweep. The rule is a substitution test: replace the literal with a second developer's value and if the sentence stays true it is CONFIGURATION and becomes `${KEY}`, and if it becomes false it is PROVENANCE and stays literal. That reverses the issue's own ranking, because all 34 hits in the densest guide file `.agents/environment.md` are provenance or named-profile definition and none is configuration, so density does not predict the defect and a blind `sed` would falsify records. Landed: three new keys `GATE_CHECKOUT`, `SHARED_STORAGE_ROOT` and `GATE_DEVICE` derived from the literals that recur and map to no existing key; the create-on-first-use route in `scripts/agent-start.py`, which printed `environment: missing` as a status label and then listed next actions that never mentioned it, so the fallback in practice was a host name copied from a document; the matching obligation in `AGENTS.md`; and one worked example. The worked example `scripts/dgx-bringup.sh` found a live defect rather than a cosmetic one: it defaulted `CUTLASS_DIR` to `$HOME/cutlass_probe` while `.agents/environment.md:389` records `$HOME/cutlass-4.5.0` as mandatory on the same box, and a configure that misses CUTLASS silently drops the sm120a NVFP4 GEMM and FlashAttention-2, which that file measures as moving the SACRED `test_qwen27_paged_engine` from 235/235 to 234/235 with the source untouched, so a stale hard-coded default is a false green. Its new test also caught that `set -a; . ./.env; set +a` over an `.env.example` copy blanks a value the caller exported, so the process-environment contract is now executable. Waves `ENV-AGNOSTIC-W1-TOOLING` through `ENV-AGNOSTIC-W5-LEDGERS` own the sweep, partition all 227 files, and open their own issues. Spec [`env-agnostic.md`](specs/env-agnostic.md) | bug | | [#1193](https://github.com/mudler/vllm.cpp/issues/1193) | `SPEC-DSPARK-QWEN3-ROUTING` | A Qwen3 DSpark draft declaring `architectures=["DSparkDraftModel"]` with `model_type` `qwen3` has no route. The pin forces every DSpark draft that is not `Qwen3DSparkModel` or `Gemma4DSparkModel` onto `model_type` `deepseek_v4` (`vllm/config/speculative.py:934-944` @ `555967922`), and vLLM PR 52197 (merged 2026-08-17 at `7075ddac`) replaced that with a leading branch normalizing the pair to `Qwen3DSparkModel`. We diverge from BOTH: the forced rewrite was never ported, so nothing in `src/vllm/entrypoints/model_loader.cpp` reads a draft config's `architectures` key at all, and `SpeculativeConfig::IsDsparkDraft` (`include/vllm/config/speculative.h:120-136`) has no production caller — every reference outside its header is in `tests/vllm/config/test_speculative_dspark.cpp:132-140`, and `ResolveSpecConfig` branches on `cli.method` alone. The checkpoint is real and gateable here: `RadixArk/Qwen3.8-27B-DSpark` at revision `85ef153be924f17ce4bf62726954eeaa4a73e854` carries exactly that config shape in one 2718576122-byte shard, drafting five layers for a 64-layer Qwen3.8-27B target | bug | | [#1213](https://github.com/mudler/vllm.cpp/issues/1213) | `ENV-LEASE-RUNTIME-STAGING` | `AGENTS.md` stated that a leased worker "has no compiler, no downloader and no Python, so it cannot produce a runtime in place", and `.agents/environment.md` carried the matching clause twice for `dgx:gpu0`. All three negatives are false. `rc describe dgx:gpu0` states that a job runs as root in an Ubuntu 24.04 container carrying `git`, `curl`, `wget`, `ssh`, `gcc`, `g++`, `make`, `cmake`, `ninja`, `pkg-config`, `python3`, `pip` and `venv`, and it instructs the reader to install anything missing; the one limit it names is the absent CUDA toolkit. Two jobs then compiled inside a lease on 2026-08-18: `claude/mudler-ubuntu-box/qwen38-gate` apt-installed `cuda-nvcc-13-0` from the `ubuntu2404/sbsa` lane and built this tree 1791/1791 to `BUILD_RC=0` (`/mnt/nas_share/rc/qwen38-gate/out-main/cfg.log` records `nvcc` 13.0.88 and `CUDA feature cutlass-fp8: ENABLED for [121a]`), and `/mnt/nas_share/rc/mtp_test/build.sh` cloned `github.com/mudler/llama.cpp` from inside a job and left a 97 MB `libggml-cuda.so` on the share. **Why it matters:** "the lease cannot produce a runtime in place" is the stated basis for treating the pinned vLLM oracle as unreachable from a lease, and that oracle is the denominator for every speed-parity number the project owes, so the premise needs re-testing rather than inheriting. This claims nothing about a model run; #1185 owns that and stays open. FIXED IN FLOW: the `AGENTS.md` paragraph and both `.agents/environment.md` clauses now say what the measurement supports, and the four real limits (no preinstalled CUDA toolkit, global installs leak until the pod restarts, CIFS `/workspace` holds no symlink so build in `/tmp` and `cp -rL`, and `-j 4` because unconstrained parallelism OOM-reboots the box) plus the host-versus-container egress distinction ride with the correction. | record | +| [#1183](https://github.com/mudler/vllm.cpp/issues/1183) | `BACKEND-ROCM` | ROCm skinny GEMM architecture eligibility caches the first device | bug | diff --git a/.agents/specs/rocm-skinny-gemm.md b/.agents/specs/rocm-skinny-gemm.md new file mode 100644 index 000000000..a21b2b064 --- /dev/null +++ b/.agents/specs/rocm-skinny-gemm.md @@ -0,0 +1,111 @@ +# ROCm wvSplitK skinny GEMM — review rework (PR #506) + +The original #506 ported upstream's `wvSplitK_hf_sml_` kernel body faithfully +(the reviewer verified the anchors and the arithmetic) but dropped the dispatch +preconditions around it. This spec records the rework; the kernel math is +unchanged. + +## The three guards, restored (all verified against the pin `555967922`) + +1. **Feature-dim guards** (upstream `utils.py:181` `m > 8 and 0 < n <= 5` with + `m = weight.shape[0]`, plus `skinny_gemms.cu:1217` `M_in % _YTILE == 0`): + with the donor's naming mapped onto ours (`out[M_tokens, N_features]`), the + kernel's YTILE=2 stores write `C[m + y + n*M_features]` unguarded for `y<2` — + on odd N the last wave writes `C[N]` (two bytes past the buffer), and at + N==1 the first wave already writes out of bounds. Our dispatch now requires + `N > 8 && (N % 2) == 0`, everything else falls through to the BLAS path. +2. **Arch guard**: the port carries only the wave32 reduction arm + (`__shfl_xor(x,16)`); upstream branches to `ROW_BCAST15/31` on gfx9 + (wave64). The gfx9 arm is NOT ported, so dispatch now refuses non-wave32 + architectures via `CapabilityFromGcnArch(DeviceArchName())` (gfx11xx/gfx12xx + only), rather than compiling and silently producing wrong sums on gfx9. +3. The `N > 8` lower bound (upstream `m > 8`) — folded into (1). + +## The test, ported for real this time + +`tests/kernels/quantization/test_rocm_skinny_gemms.py::test_rocm_wvsplitk_kernel` +@ pin — preserved: the applicable NKM factor list (tokens 1–4 = our template +arms), the xavier on/off scaling, and the **elementwise** tolerance +(`atol = eps_bf16 * sqrt(K)`, `rtol = 1e-2`; torch assert_close semantics) in +place of the aggregate NMSE. Added guard-boundary cases the upstream suite +implies: features ≤ 8 and odd features must route to BLAS and stay correct, +odd K declines, and a K%512 ≠ 0 shape exercises the K-tail. Every case runs +into a **sentinel-padded output buffer** (0xDEAD guard band) so any residual +out-of-bounds store fails the test outright. Deferred with reason recorded: +fp16 (port is bf16-only), bias (the `vt::MatmulBT` seam has no bias operand), +padded strides (our dispatch precondition is contiguous rows). + +Mutation proof: with the `N % 2` guard removed, the odd-features case corrupts +the sentinel band and the case fails; with it restored, green. + +## Boundaries + +- Kernel body unchanged from the reviewed port. +- The gfx9 (wave64) arm remains owed — a future port of the ROW_BCAST + reduction, gated the same way. +- `VT_ROCM_SKINNY=0` remains the A/B rollback; the allowlist carries it once, + in main's re-sorted layout. + +## Issue #1183 repair + +[Issue #1183](https://github.com/mudler/vllm.cpp/issues/1183) found that the +architecture guard cached the first device's architecture for the process. + +### Diagnosis + +`SkinnyGemmArchOk(int device_index)` stored `DeviceArchName(device_index)` in a +function-static string. A gfx11 call initialized that string as eligible. A +later gfx9 or unknown device then reused the gfx11 result and could reach the +wave32-only kernel. The four devices on the repair host are gfx1100, so the +test uses a controlled resolver instead of claiming heterogeneous hardware. + +### Decision + +The production call and the test now use the same HIP-free predicate. A +per-thread vector keys each result by resolver and device index. The first call +for a key resolves and parses the architecture. Later calls read one boolean +without a HIP query or a process-wide mutex. The resolver key prevents the test +resolver from contaminating a production result in the same process. + +The guard continues to accept only gfx11 and gfx12. It refuses gfx9 and every +unknown architecture. The shape, dtype, rollback, fallback, and GetBlas rules +remain unchanged. + +### Rejected alternatives + +- A single first-device value repeats the defect and is unsafe after a device + hop. +- An uncached `DeviceArchName` call adds a HIP property query to every skinny + GEMM dispatch. +- A process-wide keyed map needs synchronization on the decode path. + +### Evidence + +The test-only refactor first preserved the faulty cache. This command compiled +the production predicate and its deterministic device-hop test: + +```sh +env LD_LIBRARY_PATH=/opt/rocm/lib:/opt/rocm/lib/llvm/lib flock /home/vikash/gpu.lock cmake --build build-hip --target test_rocm_arch -j2 +``` + +The build exited 0. The next command exited 1 before the fix: + +```sh +env LD_LIBRARY_PATH=/opt/rocm/lib:/opt/rocm/lib/llvm/lib flock /home/vikash/gpu.lock build-hip/tests/test_rocm_arch '--test-case=skinny GEMM architecture eligibility follows device hops' +``` + +The gfx9 and unknown checks received `true`. Resolver counts were `{1,0,0,0}` +instead of `{1,1,1,1}`. After the keyed cache change, the same test passed 1 of +1 cases and 6 of 6 assertions with exit 0. + +The mutation replaced the keyed cache call with one function-static resolved +architecture. The build exited 0, and the same focused test exited 1 with the +same three assertion failures. After restoration, SHA-256 values for the +header, production caller, and test matched their pre-mutation values. A fresh +rebuild and focused run then passed 1 of 1 cases and 6 of 6 assertions. + +### Outcome + +Architecture eligibility follows the requested device on every device-hop +sequence. A repeated key does not query the resolver again. The gfx9 wave64 arm +remains refused and owed as recorded in `## Boundaries`. diff --git a/CMakeLists.txt b/CMakeLists.txt index f425fb5f1..fd2b4b660 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1532,6 +1532,7 @@ if(VLLM_CPP_HIP) src/vt/rocm/rocm_gdn_postconv.hip src/vt/rocm/rocm_gdn_scan.hip src/vt/rocm/rocm_gdn_fused.hip + src/vt/rocm/rocm_skinny_gemm.hip src/vt/rocm/rocm_ops.hip) if(VLLM_CPP_HIP_ARCHITECTURES) set_source_files_properties( @@ -1552,6 +1553,7 @@ if(VLLM_CPP_HIP) src/vt/rocm/rocm_gdn_postconv.hip src/vt/rocm/rocm_gdn_scan.hip src/vt/rocm/rocm_gdn_fused.hip + src/vt/rocm/rocm_skinny_gemm.hip src/vt/rocm/rocm_ops.hip PROPERTIES HIP_ARCHITECTURES "${VLLM_CPP_HIP_ARCHITECTURES}") endif() diff --git a/docs/USAGE.md b/docs/USAGE.md index c34b6d9d0..e3eb634cf 100644 --- a/docs/USAGE.md +++ b/docs/USAGE.md @@ -95,7 +95,16 @@ MoE-path coverage is partial: `MoeRouterTopK` (f32/bf16 logits, ungrouped softmax, no bias) and `MoeSiluMul` are native; the remaining chain (`kSharedExpertGate`, `kMoeCombine`/`kMoeCombineGate`, and the grouped quant expert GEMM) is not registered yet, so MoE-bearing models still throw on -those ops. On a +those ops. + +### ROCm decode GEMM routing (wvSplitK skinny path) + +Decode-shaped GEMMs (M<=4, bf16) route to a split-K skinny-GEMM kernel (a port +of vLLM's `wvSplitK`) instead of the 128x128-tile rocBLAS GEMM that dominates +decode GPU time ([#487](https://github.com/mudler/vllm.cpp/issues/487)). On by +default where it fits; `VT_ROCM_SKINNY=0` restores the BLAS path for A/B. + +On a discrete card there is no CPU fallback tier, so a model whose layers call an op that is not registered yet fails loudly with `vt: no kernel for op N on device type 5` — that is the memory-safety design working, not a crash. Run with diff --git a/include/vt/rocm/rocm_skinny_gemm_arch.h b/include/vt/rocm/rocm_skinny_gemm_arch.h new file mode 100644 index 000000000..bdc6814c4 --- /dev/null +++ b/include/vt/rocm/rocm_skinny_gemm_arch.h @@ -0,0 +1,53 @@ +// ROCm wvSplitK architecture eligibility. This header stays free of HIP +// headers so a controlled resolver can gate device-hop behavior on any host. +#pragma once + +#include +#include + +#include "vt/rocm/rocm_arch.h" + +namespace vt::rocm { + +using SkinnyGemmArchResolver = std::string (*)(int) noexcept; + +// The wvSplitK port carries only the wave32 reduction arm. Upstream uses +// ROW_BCAST15/31 on gfx9 wave64 devices, and that arm is not ported. The +// predicate accepts gfx11 and gfx12 only. An unknown architecture refuses. +namespace detail { + +class SkinnyGemmArchCache { + public: + bool Eligible(int device_index, SkinnyGemmArchResolver resolve) { + for (const Entry& entry : entries_) { + if (entry.device_index == device_index && entry.resolve == resolve) { + return entry.eligible; + } + } + + const auto cap = CapabilityFromGcnArch(resolve(device_index)); + const bool eligible = cap.has_value() && (cap->first == 11 || cap->first == 12); + entries_.push_back(Entry{device_index, resolve, eligible}); + return eligible; + } + + private: + struct Entry { + int device_index; + SkinnyGemmArchResolver resolve; + bool eligible; + }; + + std::vector entries_; +}; + +} // namespace detail + +inline bool SkinnyGemmArchOk(int device_index, SkinnyGemmArchResolver resolve) { + // Each worker reads a device property once per device. Per-thread storage + // keeps the decode path free of a process-wide lock. + static thread_local detail::SkinnyGemmArchCache cache; + return cache.Eligible(device_index, resolve); +} + +} // namespace vt::rocm diff --git a/scripts/env-doc-allowlist.txt b/scripts/env-doc-allowlist.txt index 1422e0210..ac075eddb 100644 --- a/scripts/env-doc-allowlist.txt +++ b/scripts/env-doc-allowlist.txt @@ -171,6 +171,7 @@ VT_RMSNORM_GATED_FAST VT_ROCM_GEMM_COMPUTE VT_ROCM_GEMV VT_ROCM_HIPBLASLT +VT_ROCM_SKINNY VT_SILU_FP4_FAST VT_SPEC_TRACE VT_SWIZZLE_IN_QUANT diff --git a/src/vt/rocm/rocm_matmul_hipblaslt.hip b/src/vt/rocm/rocm_matmul_hipblaslt.hip index a1d043702..305819102 100644 --- a/src/vt/rocm/rocm_matmul_hipblaslt.hip +++ b/src/vt/rocm/rocm_matmul_hipblaslt.hip @@ -28,6 +28,8 @@ #include "vt/ops.h" #include "vt/rocm/rocm_device_bind.h" #include "vt/rocm/rocm_getblas_dualslot.h" +#include "vt/rocm/rocm_runtime.h" +#include "vt/rocm/rocm_skinny_gemm_arch.h" namespace vt::rocm { namespace { @@ -154,6 +156,17 @@ bool GemvEnabled() { return on; } +// wvSplitK skinny GEMM (rocm_skinny_gemm.hip, #487): the split-K/LDS-staged +// decode path that beats the 128x128-tile rocBLAS GEMM at M<=4. Default ON for +// decode-skinny shapes; VT_ROCM_SKINNY=0 restores the BLAS path for A/B. +bool SkinnyGemmEnabled() { + static const bool on = [] { + if (const char* e = std::getenv("VT_ROCM_SKINNY")) return e[0] != '0'; + return true; + }(); + return on; +} + // y[n] = alpha * dot(x[0:K], W[n,0:K]) + beta * y[n] // grid = N — one block per output row; x cached in LDS; block-reduce over K. __global__ void Bf16GemvBTRowKernel(__hip_bfloat16* __restrict__ y, @@ -454,6 +467,11 @@ void MatmulKernelRocm(Queue& q, Tensor& out, const Tensor& a, const Tensor& b) { } // out[M,N] = a[M,K] @ b[N,K]^T +// wvSplitK skinny-GEMM host entry (rocm_skinny_gemm.hip, #487). External +// vt::rocm linkage to match the definition; declared beside its only caller. +void WvSplitKBT(hipStream_t s, void* out, const void* a, const void* b, int M, int N, + int K, int device); + // Row-major trick: gemm(OP_T, OP_N, N, M, K, B, K, A, a_rs, C, N) // BLAS: C = op(A)*op(B) with opA=T => A is KxN in col form = row B[N,K] // opB=N, B is KxM col = row A[M,K] with ld=a_rs @@ -482,6 +500,26 @@ void MatmulBTKernelRocm(Queue& q, Tensor& out, const Tensor& a, const Tensor& b) throw std::runtime_error("vt rocm: matmul_bt: bad a stride"); } + // Decode-skinny (M<=4) bf16: the wvSplitK port wins over the 128-tile GEMM + // (#487). The dispatch guards mirror the donor EXACTLY (review sweep on + // #506 found all three of these dropped, two of them memory-unsafe): + // * N > 8 (upstream `m > 8`, utils.py:181 — the feature-dim lower bound; + // at N==1 the kernel's y=1 store writes C[1] past the end of a 1-wide + // output); + // * N % YTILE == 0 (skinny_gemms.cu:1217 `M_in % _YTILE == 0`; YTILE=2 — + // the last wave's y=1 store lands at C[N] on odd N); + // * wave32 arch only (the gfx9 ROW_BCAST reduction arm is not ported); + // * K % 8 == 0 and the LDS-fit (K*M activation fits the staging buffer); + // everything else stays on the BLAS path. + if (bf16 && out.dtype == DType::kBF16 && M >= 1 && M <= 4 && (K % 8) == 0 && + N > 8 && (N % 2) == 0 && a.stride[0] == K && K * M <= 32768 && + SkinnyGemmEnabled() && + vt::rocm::SkinnyGemmArchOk(q.device.index, vt::rocm::DeviceArchName)) { + WvSplitKBT(s, out.data, a.data, b.data, static_cast(M), static_cast(N), + static_cast(K), q.device.index); + return; + } + // 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); diff --git a/src/vt/rocm/rocm_skinny_gemm.hip b/src/vt/rocm/rocm_skinny_gemm.hip new file mode 100644 index 000000000..411ea672e --- /dev/null +++ b/src/vt/rocm/rocm_skinny_gemm.hip @@ -0,0 +1,175 @@ +// ROCm skinny GEMM for decode (BACKEND-ROCM; issue #487). +// Port of vLLM's wvSplitK_hf_sml_ (csrc/rocm/skinny_gemms.cu:351-573) — the +// split-K, LDS-staged, CU-count-aware kernel that wins M<=4 shapes, replacing +// the 128x128-macro-tile rocBLAS GEMM that dominates decode GPU time. gfx1100 +// (GFX1X/wave32) bf16. De-torched; semantics mirror the donor exactly. +// +// Layout mapping to our MatmulBT (out[M,N] = a[M,K] @ b[N,K]^T, decode M<=4): +// donor in_a (weight [N,K]) = our b; donor in_b (x [M,K]) = our a; +// donor M_in = our N (output dim); donor N_in = our M; C[N_in,M_in] = out[M,N]. + +#include +#include + +#include +#include +#include + +#include "vt/ops.h" + +namespace vt::rocm { +namespace { + +constexpr int kLdsSize = 64 * 1024; // gfx1100 (non-gfx95x) +constexpr int kThrds = 32; // wave32 +constexpr int kWvPrGrp = 16; +constexpr int kYtile = 2; // proven decode config (donor gfx1x, N small) +constexpr int kUnrl = 2; +constexpr int kAChunk = 8; + +using scalar8 = __attribute__((__vector_size__(4 * sizeof(float)))) float; +union bigType { + __hip_bfloat16 h[kAChunk]; + float f[kAChunk / 2]; + scalar8 h8; +}; + +__device__ __forceinline__ unsigned int min__(uint32_t a, uint32_t b) { + return min(a, b); +} + +// Donor's persistent-workgroup splitter (host-side in skinny_gemms.cu:1169). +inline int mindiv(int N, int div1, int div2) { + int nPrRnd = div1 * div2; + int rnds[13]; + for (int i = 0; i < 13; i++) { + rnds[i] = (N + nPrRnd - 1) / nPrRnd; + nPrRnd -= div1; + } + for (int i = 12; i >= 0; i--) + if (rnds[0] == rnds[i]) return (div2 - i); + return 0; +} + +// A (activation) fits LDS. N = decode batch (our M). bf16, f32 accum. +template +__global__ void __launch_bounds__(kWvPrGrp * kThrds) + wvSplitKSml(const int K, const int Kbp, const int Kap, const int M, + const __hip_bfloat16* __restrict__ B, const __hip_bfloat16* __restrict__ A, + __hip_bfloat16* C, const int _WvPrGrp, const int CuCount) { + constexpr int max_lds_len = kLdsSize / 2; // bf16 elements + __shared__ __hip_bfloat16 s[max_lds_len]; + + // Stage the activation row(s) A [N,K] into LDS. + for (uint32_t k = (threadIdx.y * kThrds + threadIdx.x) * kAChunk; + k < min__(Kap * N, max_lds_len); k += kThrds * kWvPrGrp * kAChunk) { + *((bigType*)(&s[k])) = *((const bigType*)(&A[k])); + } + __syncthreads(); + + if (threadIdx.y >= _WvPrGrp) return; + + uint32_t m = (blockIdx.x * _WvPrGrp + (threadIdx.y % _WvPrGrp)) * kYtile; + + while (m < static_cast(M)) { + float sum[N][kYtile] = {}; + for (uint32_t k1 = 0; k1 < static_cast(K); k1 += kThrds * kAChunk * kUnrl) { + bigType bigA[N][kUnrl] = {}; + bigType bigB[kYtile][kUnrl]; +#pragma unroll + for (uint32_t k2 = 0; k2 < kUnrl; k2++) { + uint32_t k = k1 + k2 * kThrds * kAChunk; + uint32_t k_ = k + threadIdx.x * kAChunk; + const __hip_bfloat16* B_ = &B[min__(k_, K - kAChunk)]; + for (int y = 0; y < kYtile; y++) + bigB[y][k2].h8 = __builtin_nontemporal_load( + (const scalar8*)(&B_[min__(y + m, M - 1) * Kbp])); + } +#pragma unroll + for (uint32_t k2 = 0; k2 < kUnrl; k2++) { + uint32_t k = k1 + k2 * kThrds * kAChunk; + uint32_t k_ = k + threadIdx.x * kAChunk; + if (k_ >= static_cast(K)) break; + for (int n = 0; n < N; n++) bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n]))); + } + // Interleaved MAC; bf16 pairs unpacked to f32 (donor DOT2C bf16 branch). + for (uint32_t k2 = 0; k2 < kUnrl; k2++) { + for (int n = 0; n < N; n++) { + for (int y = 0; y < kYtile; y++) { +#pragma unroll + for (uint32_t b = 0; b < kAChunk / 2; b++) { + float2 a2 = __bfloat1622float2(*((__hip_bfloat162*)(&(bigA[n][k2].h[b * 2])))); + float2 b2 = __bfloat1622float2(*((__hip_bfloat162*)(&(bigB[y][k2].h[b * 2])))); + sum[n][y] += (a2.x * b2.x) + (a2.y * b2.y); + } + } + } + } + } + __builtin_amdgcn_sched_barrier(0); + // Wave32 reduction: DPP row_shr 8/4/2/1 then shfl_xor(16). + for (int n = 0; n < N; n++) { + for (int y = 0; y < kYtile; y++) { + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x118, 0xf, 0xf, 1); + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x114, 0xf, 0xf, 1); + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x112, 0xf, 0xf, 1); + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x111, 0xf, 0xf, 1); + sum[n][y] += __shfl_xor(sum[n][y], 16); + } + } + if (threadIdx.x == (kThrds - 1)) { + for (int n = 0; n < N; n++) + for (int y = 0; y < kYtile; y++) C[m + y + n * M] = __float2bfloat16(sum[n][y]); + } + m += CuCount * _WvPrGrp * kYtile; + } +} + +int DeviceCuCount(int device) { + static const int cache = [] { + int dev = 0; + if (hipGetDevice(&dev) != hipSuccess) return 96; // gfx1100 fallback + int n = 0; + if (hipDeviceGetAttribute(&n, hipDeviceAttributeMultiprocessorCount, dev) != hipSuccess || + n <= 0) + return 96; + return n; + }(); + (void)device; + return cache; +} + +} // namespace + +// out[M,N] = a[M,K] @ b[N,K]^T, bf16 in/out. Only called for the decode-skinny +// gate below (M in 1..4, K%8==0); the caller keeps every other shape on the +// BLAS path. +void WvSplitKBT(hipStream_t s, void* out, const void* a, const void* b, int M, int N, + int K, int device) { + const int cu = DeviceCuCount(device); + dim3 grid(cu), block(kThrds, kWvPrGrp); + const int wvPrGrp = mindiv(N, cu * kYtile, kWvPrGrp); + auto* C = static_cast<__hip_bfloat16*>(out); + auto* A = static_cast(a); + auto* B = static_cast(b); + switch (M) { + case 1: + wvSplitKSml<1><<>>(K, K, K, N, B, A, C, wvPrGrp, cu); + break; + case 2: + wvSplitKSml<2><<>>(K, K, K, N, B, A, C, wvPrGrp, cu); + break; + case 3: + wvSplitKSml<3><<>>(K, K, K, N, B, A, C, wvPrGrp, cu); + break; + case 4: + wvSplitKSml<4><<>>(K, K, K, N, B, A, C, wvPrGrp, cu); + break; + default: + throw std::runtime_error("vt rocm: wvSplitK unsupported M=" + std::to_string(M)); + } + if (hipGetLastError() != hipSuccess) + throw std::runtime_error("vt rocm: wvSplitK launch failed"); +} + +} // namespace vt::rocm diff --git a/tests/vt/test_backend_cross_device.cpp b/tests/vt/test_backend_cross_device.cpp index fc5d0e61b..87400aeed 100644 --- a/tests/vt/test_backend_cross_device.cpp +++ b/tests/vt/test_backend_cross_device.cpp @@ -2073,6 +2073,104 @@ TEST_CASE("MoeRouterTopK matches the CPU oracle (f32 and bf16 logits)") { } } +TEST_CASE("decode-skinny MatmulBT (wvSplitK path) matches the CPU oracle") { + // Port of upstream's tests/kernels/quantization/test_rocm_skinny_gemms.py + // ::test_rocm_wvsplitk_kernel @ pin 55596792 (review sweep on #506: the first + // version of this case had aggregate-NMSE tolerance that ten completely + // wrong elements would still pass, every K a multiple of the 512 stride so + // the K-tail path never ran, and no guard-boundary shapes at all). + // + // Preserved from upstream: the NKM factor list (the applicable subset — see + // below), the xavier on/off scaling, and the ELEMENTWISE tolerance + // atol = eps_bf16 * sqrt(K), rtol = 1e-2 (torch.testing.assert_close + // semantics). Deferred with reason: fp16 (our port is bf16-only), bias + // (the vt::MatmulBT seam has no bias operand), padded strides (our dispatch + // requires contiguous rows — a documented precondition), and the fp8/rc + // kernel variants (not ported). The (n,k,m) upstream triple = (tokens, K, + // features) here. + struct Shape { int64_t tok, k, feat; const char* why; }; + const Shape shapes[] = { + // the upstream sweep (token counts 1-4 = our template arms) + {1, 32, 16, "upstream"}, {1, 64, 64, "upstream"}, {2, 256, 256, "upstream"}, + {3, 1024, 1024, "upstream"}, {4, 4096, 4096, "upstream"}, + // K-tail: K % 512 != 0 exercises the `if (k_ >= K) break` remainder path + {4, 4096 + 16, 4096, "k-tail"}, {1, 9216, 512, "upstream"}, + // guard boundaries (must stay CORRECT via the BLAS fallback) + {2, 256, 8, "features<=8 declines (upstream m>8)"}, + {2, 256, 254, "even below bound: takes skinny"}, + {2, 256, 255, "odd features decline (YTILE=2 OOB class)"}, + {2, 254, 256, "K%8!=0 declines"}, + }; + const double kEpsBf16 = 0.0078125; // 2^-8 + for (const Shape& sh : shapes) { + for (bool xnorm : {false, true}) { + CAPTURE(sh.why); + CAPTURE(sh.tok); + CAPTURE(sh.k); + CAPTURE(sh.feat); + CAPTURE(xnorm); + const int64_t M = sh.tok, N = sh.feat, K = sh.k; + const size_t an = static_cast(M) * K, bn = static_cast(N) * K; + const double xavier = xnorm ? std::sqrt(2.0 / static_cast(K)) : 1.0; + std::vector a = RandomVec(an, 991, -1.0f, 1.0f); + std::vector b = RandomVec(bn, 992, -1.0f, 1.0f); + for (float& x : a) x = static_cast(x * xavier); + for (float& x : b) x = static_cast(x * xavier); + const std::vector a_bf = Bf16Bits(a), b_bf = Bf16Bits(b); + + std::vector ref(static_cast(M) * N, 0); + { + vt::Backend& cpu = vt::GetBackend(DeviceType::kCPU); + Queue cq = cpu.CreateQueue(); + const Device cd{DeviceType::kCPU, 0}; + std::vector ca = a_bf, cb = b_bf; + Tensor ta = Tensor::Contiguous(ca.data(), DType::kBF16, cd, {M, K}); + Tensor tb = Tensor::Contiguous(cb.data(), DType::kBF16, cd, {N, K}); + Tensor to = Tensor::Contiguous(ref.data(), DType::kBF16, cd, {M, N}); + vt::MatmulBT(cq, to, ta, tb); + cpu.DestroyQueue(cq); + } + for (DeviceType dt : RegisteredDevices()) { + if (!OpAvailable(vt::OpId::kMatmulBT, dt)) continue; + CAPTURE(DeviceName(dt)); + vt::Backend& dev = vt::GetBackend(dt); + Queue q = dev.CreateQueue(); + const Device d{dt, 0}; + // Sentinel-padded output: the dispatch must never write past M*N + // elements (the odd-features OOB class from the review). + const size_t out_elems = static_cast(M) * N; + const size_t guard_elems = 128; + DevBufBytes da(dev, q, an * 2), db(dev, q, bn * 2), + dout(dev, q, (out_elems + guard_elems) * 2); + std::vector fill(out_elems + guard_elems, 0xDEAD); + dout.Upload(fill.data()); + da.Upload(a_bf.data()); + db.Upload(b_bf.data()); + Tensor ta = Tensor::Contiguous(da.ptr(), DType::kBF16, d, {M, K}); + Tensor tb = Tensor::Contiguous(db.ptr(), DType::kBF16, d, {N, K}); + Tensor to = Tensor::Contiguous(dout.ptr(), DType::kBF16, d, {M, N}); + vt::MatmulBT(q, to, ta, tb); + std::vector got(out_elems + guard_elems); + dout.Download(got.data()); + // Elementwise tolerance (upstream assert_close), never aggregate NMSE. + const double atol = kEpsBf16 * std::sqrt(static_cast(K)); + for (size_t i = 0; i < out_elems; ++i) { + uint32_t ug = static_cast(got[i]) << 16, ur = static_cast(ref[i]) << 16; + float gf, rf; + std::memcpy(&gf, &ug, 4); + std::memcpy(&rf, &ur, 4); + CHECK(std::fabs(gf - rf) <= atol + 1e-2 * std::fabs(rf)); + } + // The guard band must be untouched by ANY path (skinny or BLAS). + for (size_t i = out_elems; i < out_elems + guard_elems; ++i) + CHECK(got[i] == 0xDEAD); + dev.DestroyQueue(q); + } + } + } +} + + TEST_CASE("reference tier: an op with no native kernel matches the CPU oracle (unified only)") { constexpr int64_t kRows = 7, kCols = 48; constexpr size_t kN = kRows * kCols; diff --git a/tests/vt/test_rocm_arch.cpp b/tests/vt/test_rocm_arch.cpp index 6cd8b7363..b916808e3 100644 --- a/tests/vt/test_rocm_arch.cpp +++ b/tests/vt/test_rocm_arch.cpp @@ -10,13 +10,31 @@ // The cases are upstream's own worked examples, from the docstring of // vllm/platforms/rocm.py:223-291 `_capability_from_gcn_arch`, plus the three // boards offered on issue #41 (gfx1100, gfx1103, gfx1151). +#include +#include + #include #include "vt/rocm/rocm_arch.h" +#include "vt/rocm/rocm_skinny_gemm_arch.h" using vt::rocm::CapabilityFromGcnArch; namespace { +std::array skinny_arch_resolves{}; + +std::string SimulatedSkinnyArch(int device_index) noexcept { + if (device_index >= 0 && device_index < static_cast(skinny_arch_resolves.size())) { + ++skinny_arch_resolves[static_cast(device_index)]; + } + switch (device_index) { + case 0: return "gfx1100"; + case 1: return "gfx942:sramecc+:xnack-"; + case 2: return "gfx1201"; + default: return "future-arch"; + } +} + // Reads as (major, minor) at the call site instead of .first / .second. void CheckArch(const char* gcn, int major, int minor) { const auto cap = CapabilityFromGcnArch(gcn); @@ -79,3 +97,13 @@ TEST_CASE("the parse is constexpr, so a wrong answer is a compile error") { static_assert(!CapabilityFromGcnArch("sm_121a").has_value()); CHECK(true); } + +TEST_CASE("skinny GEMM architecture eligibility follows device hops") { + skinny_arch_resolves.fill(0); + CHECK(vt::rocm::SkinnyGemmArchOk(0, SimulatedSkinnyArch)); + CHECK_FALSE(vt::rocm::SkinnyGemmArchOk(1, SimulatedSkinnyArch)); + CHECK(vt::rocm::SkinnyGemmArchOk(2, SimulatedSkinnyArch)); + CHECK_FALSE(vt::rocm::SkinnyGemmArchOk(3, SimulatedSkinnyArch)); + CHECK(vt::rocm::SkinnyGemmArchOk(0, SimulatedSkinnyArch)); + CHECK(skinny_arch_resolves == std::array{1, 1, 1, 1}); +}