diff --git a/.agents/specs/rocm-gg-keep-quant.md b/.agents/specs/rocm-gg-keep-quant.md new file mode 100644 index 000000000..30f6aaa6b --- /dev/null +++ b/.agents/specs/rocm-gg-keep-quant.md @@ -0,0 +1,58 @@ +# ROCm keep-quant expert GEMM — review rework (PR #523) + +## What this fixes + +The review sweep (localai-bot, 2026-08-13) found the original #523 shape +registered `kMatmulBTQuant` with a loader that flips keep-quant on a BOOLEAN +(`GgufQuantComputeAvailable()` = `OpRegistered(...)`), while the ROCm kernel +implements 4 of the 12 formats the loader admits (Q4_0, Q8_0, Q2_K, Q3_K, Q4_K, +Q5_K, Q6_K, IQ2_XXS, IQ3_XXS, IQ2_S, MXFP4). On a discrete card with no CPU +fallback tier, a Q4_0/Q2_K/IQ2 model that loaded and generated fine before +would keep blocks quantized and throw at first forward. Same boolean flipped +`keep_f16` on, and the ROCm `MatmulBT` refuses f16 — a second regression of a +working path. + +## The rework + +1. **Per-dtype capability in the loader** (`gguf_keep_quant.cpp`): + `KeepQuantDType` and the keep-f16 default now consult the running device's + actual support. ROCm keep-quant supports {Q8_0, Q4_K, Q5_K, Q6_K} + (kMatmulBTQuant + kMatmulBTQuantGrouped both); ROCm keep-f16 is OFF + (`MatmulBTKernelRocm` accepts bf16/f32 only). Unsupported formats keep the + pre-existing `expand_bf16` residency — no load fails, no forward throws, and + `VT_GGUF_KEEP_QUANT=1` on a Q4_0 model is a no-op rather than a regression. + CUDA/CPU behavior is byte-identical (their sets already cover the CPU list). +2. **Capture-safe scratch**: the per-call `hipMalloc`/`hipFree`/ + `hipStreamSynchronize` on the activation-quant scratch (illegal under + hipGraph stream capture — blocks #473/#332) becomes a grow-only per-stream + pool via `hipMallocAsync`, mirroring the donor's `EnsureScratch` + + `RetireGraphScratch` (never-freed, because a captured graph may have baked + the pointer). Also fixes the `qact` leak when `Check()` threw between + malloc and free. +3. **The refusal messages** name the actually-unported formats (Q4_0, Q2_K, + Q3_K, IQ2_XXS, IQ3_XXS, IQ2_S, MXFP4) instead of double-listing Q5_K — the + message is now unreachable in practice (the loader pre-filters) but stays + correct as the last line of defense. +4. **Teeth**: the non-grouped `kMatmulBTQuant` gains its own cross-device case + (it carried the headline mechanism and had no test), and both new cases + `REQUIRE(OpAvailable(...))` instead of skipping silently when registration + is dropped. The grouped case keeps its NMSE<=5e-4 vs CPU keep-quant oracle + bar. +5. `Dp4a` keeps the portable four-MAC body if `__dp4a` is absent on the + gfx1100 toolchain (verified at build time); if `__dp4a` compiles, use it. + +## Gates + +- Focused: `test_backend_cross_device` (grouped + non-grouped keep-quant + cases, REQUIRE-proven registration), red-first by stash-revert. +- Regression: the 0.8B + 0.6B M4 gates; Qwen3.6-35B-A3B Q4_K_M e2e on one + gfx1100 card (`--max-num-seqs 1`); a Q4_0 GGUF load on ROCm proving no + regression (expands, generates, no throw). +- Full HIP ctest zero-delta vs base. + +## Boundaries + +- No change to the ported dot-product cores (review verified them against the + donor, DotQ6K byte-for-byte). +- Q2_K/Q3_K/IQ2/IQ3/MXFP4 ROCm kernels remain owed and are recorded as such + here and in the refusal messages. diff --git a/.agents/specs/rocm-grouped-quant-gemm.md b/.agents/specs/rocm-grouped-quant-gemm.md new file mode 100644 index 000000000..7a555f3d7 --- /dev/null +++ b/.agents/specs/rocm-grouped-quant-gemm.md @@ -0,0 +1,77 @@ +# ROCm grouped quant expert GEMM (kMatmulBTQuantGrouped) — spike + +**Issue:** #41 (ROCm lane); the named blocker for MoE-bearing models after the +GDN slice (#334–#345) and the MoE chain (#348, #509). +**Status:** spike — no code yet. + +## The gap, verified + +On discrete ROCm the MoE path now resolves everything except the expert GEMM: +`kMoeRouterTopK` + `kMoeSiluMul` (#348), `kSharedExpertGate`/`kMoeCombine`/ +`kMoeCombineGate` (#509) are native and gated. The remaining throw is +`kMatmulBTQuantGrouped` — the keep-quant grouped expert GEMM that runs the +stacked `[E*N,K]` expert towers. Without it, MoE-bearing models +(Qwen3.5-27B-class GDN-MoE, DeepSeek-V4 GGUF) throw on discrete ROCm. + +`test_bench`/`test_capi` flipped green once the chain ops landed (they don't +reach the grouped GEMM). `test_loaded_engine_dense` still fails — but on the +**async-scheduling assertion** (`runner_supports_async()=false` on ROCm), a +lane capability gap unrelated to this op. + +## What the donor actually is + +`src/vt/cuda/cuda_quant_dot.cu` (2069 lines). The grouped GEMM is +`QuantDotGemmGroupedKernel` (:746) + a fused SwiGLU variant (:799) + +Q8_0-specific kernels (:1404/:1441). Structure: + +1. **Shared activation quant**: input rows quantized to Q8_K once + (`QuantizeRowQ8_K`, CPU ref `cpu_quant_act.cpp:88`; a `QuantizeQ8_0Kernel` + device quantizer exists for the Q8_0 path). +2. **Per-format integer dot superblocks**: `DotSuperblock` specializations + (`:655`+) for Q2_K, Q3_K, Q4_K, Q5_K, Q6_K, Q4_0, Q8_0, IQ2_XXS/IQ3_XXS — + each dequantizes a keep-quant weight superblock and dots against the Q8_K + activation block. This is the bulk and the only genuinely tricky part. +3. **Grouped dispatch**: warp-per-(p,j), `__shfl_down_sync` reduction + (HIP-compatible as-is), expert row selected by `expert_ids[p]`. + +The CPU reference (`cpu_quant_dot.cpp` VecDot family) is complete and is the +gate oracle. HIP needs no torch; the donor's torch surface is only the host +glue. + +## Port plan (per-format PRs, red-first, CPU-oracle gated) + +- **W0: Q8_K activation quant + Q8_0 dot + grouped skeleton.** Smallest + end-to-end slice that runs a real (if low-value) grouped GEMM; establishes + the registration, the Q8_K quantizer port, and the cross-device gate vs + `VecDotQ8_0Q8_0` (cpu_quant_dot.cpp:88). RED: op unregistered today. +- **W1: Q4_0 + Q4_K** (`VecDotQ4_0Q8_0` :50, `VecDotQ4_KQ8_K` :203) — the + dominant GGUF expert formats. +- **W2: Q5_K/Q6_K/Q2_K** and the fused SwiGLU variant (the ds4 epilogue). +- **W3: IQ2/IQ3** — lowest-value, last. + +Each family: hand-port from the donor's `DotSuperblock`, cross-device case vs +the CPU VecDot oracle (NMSE ≤ 5e-4 — the same band the CUDA lane uses, since +the integer core is bit-exact and only the float scale sum reassociates), +focused + full gate. + +## Testability constraint (honest) + +The op is only reachable end-to-end on MoE models I cannot fit on this box +(Qwen3.5-27B needs a multi-GB GGUF; the 0.8B has no experts). So the gate is +the **CPU reference at the op level** (cross-device, both groupings, the +broadcast-activation arm), and the model-level e2e stays PENDING a host with +the checkpoint — that is a real constraint, stated, not papered over. + +## What is deliberately not in scope + +- The fused SwiGLU grouped kernel (W2, a perf/composition variant). +- ggml's SIMD-table IQ formats' fastest paths (port the reference math first). +- Any perf tuning — correctness first; the win over "no path at all" is + binary. + +## Stop conditions + +- A format's dot cannot be made NMSE-clean vs the CPU VecDot oracle → stop and + post the failing evidence on #41 rather than ship a wrong quant path. +- A model-level e2e claim is ever made from op-level-only evidence → it must + not be; the constraint above holds. diff --git a/CMakeLists.txt b/CMakeLists.txt index a31be7aae..901e4e75b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1421,6 +1421,8 @@ if(VLLM_CPP_HIP) src/vt/rocm/rocm_gemma4_expert_geglu.hip src/vt/rocm/rocm_fp8_channel_gemv.hip src/vt/rocm/rocm_moe_router.hip + src/vt/rocm/rocm_moe_chain.hip + src/vt/rocm/rocm_grouped_gemm.hip src/vt/rocm/rocm_sample.hip src/vt/rocm/rocm_gdn_state.hip src/vt/rocm/rocm_gdn_conv.hip @@ -1441,6 +1443,8 @@ if(VLLM_CPP_HIP) src/vt/rocm/rocm_gemma4_expert_geglu.hip src/vt/rocm/rocm_fp8_channel_gemv.hip src/vt/rocm/rocm_moe_router.hip + src/vt/rocm/rocm_moe_chain.hip + src/vt/rocm/rocm_grouped_gemm.hip src/vt/rocm/rocm_sample.hip src/vt/rocm/rocm_gdn_state.hip src/vt/rocm/rocm_gdn_conv.hip diff --git a/docs/FEATURES.md b/docs/FEATURES.md index 26317660c..5b22fc801 100644 --- a/docs/FEATURES.md +++ b/docs/FEATURES.md @@ -242,7 +242,7 @@ the only ☐ in our column here. | CPU (x86, Arm i8mm; A76 assembly correct/default, llama speed gate open) | ✅ | ◐ | ☐ | ✅ | | Metal (Apple Silicon) | ✅ | ☐ | ☐ | ✅ | | Vulkan | ◐ | ☐ | ☐ | ✅ | -| ROCm | W0 verified on 5 gfx archs; dense and GDN models run all-native. Strict CPU parity is open in the measured near-tie regime (#269) | 44 registered ops including full GDN; ctest-green gfx1151/1103/1100/1201/1200 ([#41](https://github.com/mudler/vllm.cpp/issues/41)). APU managed allocation is unverified. [ROCM.md](ROCM.md) | ✅ | ✅ | +| ROCm | W0 verified on 5 gfx archs; dense and GDN models run all-native. Strict CPU parity is open in the measured near-tie regime (#269) | 47 registered ops including full GDN and MoE combine/gate; ctest-green gfx1151/1103/1100/1201/1200 ([#41](https://github.com/mudler/vllm.cpp/issues/41)). APU managed allocation is unverified. [ROCM.md](ROCM.md) | ✅ | ✅ | | XPU / TPU | ☐ | ✅ | ◐ | ☐ | | Tenstorrent Blackhole | ◐ `ACTIVE`, OPT-125m STRICT 6/6 e2e; Qwen3-0.6B gate wired with device goldens. Full 16x16 rerun and residual-RMS numerics at the rows≥32 device boundary both owed ([spec](../.agents/specs/tenstorrent-backend.md)) | ✅ | ☐ | ☐ | @@ -330,7 +330,7 @@ CPU elementwise GEMM (f32/f16/bf16) runs AVX2 and AVX-512 tiers on x86 where the | LoRA end to end | CPU brick landed | Unwired standalone; not usable through the server | | Multimodal over HTTP | Image request path wired; forward + codec pending | `ROAD-V1-MM` W1-W3 landed (`server_main.cpp:826`). Open: no mm-forward consuming `Request.mm_features`; no image codec vendored (raw RGB only); video/audio/multi-image not started | | Reranking / classify models | Engine side only | Embeddings are LIVE (`LlamaModel`, `vllm_embed`, `/v1/embeddings`); the classify/score heads are landed ops with no registered arch | -| ROCm | W0 community-verified on 5 gfx archs; classic-dense and GDN-hybrid e2e run all-native; correctness gaps remain | 44 registered ops including the GDN state/conv/postconv/recurrence set; APU managed-allocation branch remains unverified. [ROCM.md](ROCM.md) | +| ROCm | W0 community-verified on 5 gfx archs; classic-dense and GDN-hybrid e2e run all-native; correctness gaps remain | 47 registered ops including the GDN state/conv/postconv/recurrence set and MoE combine/gate; APU managed-allocation branch remains unverified. [ROCM.md](ROCM.md) | | XPU, TPU | Not started | CUDA, CPU, Metal and Vulkan are the built backends | | Custom logits processors on CUDA | Open, not root-caused | Segfaults in a CUDA build, 232/232 green on CPU | | Memory budgeting (`ROAD-V1-MEM`, #83) | M1+M2 landed (absolute bytes) | `--kv-cache-memory` sizes the KV pool from an absolute byte budget (ABI v16, group-aware divisor); `--num-blocks` overrides; `--gpu-memory-utilization` needs the M3 profile run (dgx-gated). See `specs/kv-sizing.md` | diff --git a/docs/USAGE.md b/docs/USAGE.md index 4f80a43ef..1adfacf4f 100644 --- a/docs/USAGE.md +++ b/docs/USAGE.md @@ -66,11 +66,12 @@ portable scan), and the norm-gate/preamble ops (`kRmsNormGated`, GDN-hybrid models call. Compressed conv/SSM state (bf16, the vLLM `mamba_cache_dtype` default) is advertised via the `SupportsCompressedConvState`/`SupportsCompressedGdnState` backend probes. -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 +MoE-path coverage: `MoeRouterTopK` (f32/bf16 logits, ungrouped softmax, no +bias), `MoeSiluMul`, `SharedExpertGate`, `MoeCombine`, and `MoeCombineGate` +are native. All three combine/gate ops accept f32 and bf16 operands and +refuse anything else with a named message (f16 is not a supported arm). The grouped quant expert GEMM (`kMatmulBTQuantGrouped`) is not +registered yet, so models whose experts run keep-quant grouped GEMMs still +throw there. 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/src/vllm/model_executor/model_loader/gguf_keep_quant.cpp b/src/vllm/model_executor/model_loader/gguf_keep_quant.cpp index aa96cc8fd..43cbcad6b 100644 --- a/src/vllm/model_executor/model_loader/gguf_keep_quant.cpp +++ b/src/vllm/model_executor/model_loader/gguf_keep_quant.cpp @@ -109,6 +109,35 @@ bool KeepF16DType(uint32_t ggml_type) { return ggml_type == 1; } // ggml type id 40 is the NVFP4 fork extension; see gguf_dequant.cpp case 40. bool KeepNvfp4DType(uint32_t ggml_type) { return ggml_type == 40; } +// Device-side keep-quant capability (review sweep on #523): the master +// boolean `GgufQuantComputeAvailable()` only says the OP is registered; a +// device's kernel set can be narrower than the CPU admission list, and on a +// discrete backend with no CPU fallback tier a format the device cannot +// execute must keep its pre-existing expand-bf16 residency -- flipping it to +// a keep-quant block throws at FORWARD time with the whole model resident. +// Per-device sets name what the registered kernels actually implement. +bool DeviceKeepQuantSupported(vt::DType dt, vt::DeviceType dev) { + switch (dev) { + case vt::DeviceType::kROCM: + // src/vt/rocm/rocm_grouped_gemm.hip implements exactly these on both the + // grouped and non-grouped arms; Q4_0/Q2_K/Q3_K/IQ2_*/IQ3_*/MXFP4 are + // owed (recorded in .agents/specs/rocm-gg-keep-quant.md). + return dt == vt::DType::kQ8_0 || dt == vt::DType::kQ4_K || + dt == vt::DType::kQ5_K || dt == vt::DType::kQ6_K; + default: + // CUDA falls back to the CPU kernel for anything it lacks + // (cuda_quant_dot.cu:1841-1846); the CPU list IS the CPU capability. + return true; + } +} + +// keep-f16 needs an f16-capable MatmulBT on the running device; the ROCm +// kernel accepts bf16/bf16 and f32/f32 only, so an F16 file weight must +// expand there rather than be kept and refused at first forward (same review). +bool DeviceKeepF16Supported(vt::DeviceType dev) { + return dev != vt::DeviceType::kROCM; +} + bool KeepQuantDType(uint32_t ggml_type, vt::DType* out) { vt::DType dt = vt::DType::kF32; if (!vt::BlockDTypeFromGgmlTypeId(ggml_type, &dt)) return false; @@ -141,9 +170,14 @@ GgufResidency RouteGgufTensor(bool keep_quant, bool keep_f16, bool nvfp4_fp4, const int64_t k = KeepQuantKDim(role, shape); vt::DType dt = vt::DType::kF32; // ggml_row_size's precondition: a row is a whole number of blocks. A weight - // whose K is ragged cannot be dotted block-wise, so it expands. + // whose K is ragged cannot be dotted block-wise, so it expands. The device + // gate (review #523): a format the RUNNING device cannot execute keeps its + // pre-existing expand-bf16 residency instead of flipping to a keep-quant + // block that throws at forward time on a card with no CPU fallback tier. if (k > 0 && KeepQuantDType(ggml_type, &dt) && - k % vt::BlockElems(dt) == 0) { + k % vt::BlockElems(dt) == 0 && + DeviceKeepQuantSupported( + dt, vllm::platforms::CurrentPlatform().device_type())) { return GgufResidency::kKeepQuant; } } @@ -190,7 +224,9 @@ GgufLoadPolicy GgufLoadPolicy::FromEnv() { // greedy tokens byte-identical (native-f16 compute, md5 d235db1... unchanged). // VT_GGUF_KEEP_F16=0 is the opt-out; rides expand_nk so it is CPU-only and off // under VT_CPU_REF regardless (the oracle load stays byte-identical). - p.keep_f16 = EnvOnOr("VT_GGUF_KEEP_F16", p.expand_nk) && p.expand_nk; + p.keep_f16 = EnvOnOr("VT_GGUF_KEEP_F16", p.expand_nk) && p.expand_nk && + DeviceKeepF16Supported( + vllm::platforms::CurrentPlatform().device_type()); // `QUANT-GGUF-NVFP4` column C. Same shape as the keep-quant default: ON // wherever the running device can execute the NVFP4 GEMM (CUDA today; a CPU // build keeps expanding, which is correct but unquantized), with diff --git a/src/vt/rocm/rocm_grouped_gemm.hip b/src/vt/rocm/rocm_grouped_gemm.hip new file mode 100644 index 000000000..289e6141b --- /dev/null +++ b/src/vt/rocm/rocm_grouped_gemm.hip @@ -0,0 +1,569 @@ +// ROCm grouped quant expert GEMM (BACKEND-ROCM; issue #41, the MoE-path +// blocker). Port of src/vt/cuda/cuda_quant_dot.cu grouped path: +// QuantizeQ8KKernel (Q8_K activation for K-quant formats) +// QuantizeQ8_0Kernel (Q8_0 activation for the Q8_0 format) +// QuantDotGemmGroupedKernel (:746) + QuantDotGemmGroupedQ8_0Kernel (:1404) +// Dot superblocks DotQ8_0 / DotQ4K / DotQ6K ported 1:1. Bit-exact integer +// cores (__dp4a); float scale products reassociate across lanes as the donor's. +// +// out[P,N] (f32/bf16) = per (p,j): sum_sb dot(Q8 act[p], keepquant w[e,j]), +// e = expert_ids[p]; activation quantized once (broadcast when 1 row). +// +// Covers the formats the target GDN-MoE GGUFs use (Q4_K / Q6_K / Q8_0). The +// IQ2/IQ3/Q2_K/Q3_K/Q5_K superblocks port identically against this skeleton. + +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "vt/ops.h" +#include "vt/rocm/rocm_device_bind.h" + +// Block layouts — the single source of truth (ggml-common.h mirrors). +#include "vt/cpu/cpu_quant_blocks.h" + +namespace vt::rocm { +namespace { + +using vt::cpu::BlockQ8_0; +using vt::cpu::BlockQ8_K; +using vt::cpu::BlockQ4_K; +using vt::cpu::BlockQ5_K; +using vt::cpu::BlockQ6_K; +using vt::cpu::kQK8_0; +using vt::cpu::kQK_K; + +enum class ActDT : int { kF32 = 0, kF16 = 1, kBF16 = 2 }; +inline ActDT ActDtOf(DType dt) { + return dt == DType::kF32 ? ActDT::kF32 : dt == DType::kF16 ? ActDT::kF16 : ActDT::kBF16; +} + +// ---- device numeric helpers (bit-exact ports from cuda_quant_dot.cu) ---- +__device__ inline float DF16ToF32(uint16_t h) { + uint32_t sign = static_cast(h & 0x8000) << 16; + uint32_t exp = (h >> 10) & 0x1F; + uint32_t mant = h & 0x3FF; + if (exp == 0x1F) return __int_as_float(sign | 0x7F800000 | (mant << 13)); + if (exp == 0) { + if (mant == 0) return __int_as_float(sign); + int shift = 0; + while ((mant & 0x400) == 0) { mant <<= 1; ++shift; } + mant &= 0x3FF; + return __int_as_float(sign | ((113 - shift) << 23) | (mant << 13)); + } + return __int_as_float(sign | ((exp + 112) << 23) | (mant << 13)); +} +__device__ inline float DBF16ToF32(uint16_t b) { + return __int_as_float(static_cast(b) << 16); +} +__device__ inline uint16_t DF32ToBF16(float f) { + uint32_t u = __float_as_int(f); + if ((u & 0x7F800000) == 0x7F800000 && (u & 0x7FFFFF)) + return static_cast((u >> 16) | 0x0040); + uint32_t rounding = 0x7FFF + ((u >> 16) & 1); + return static_cast((u + rounding) >> 16); +} +__device__ inline uint16_t DF32ToF16(float f) { + uint32_t u = __float_as_uint(f); + uint16_t sign = static_cast((u >> 16) & 0x8000); + int32_t exp = static_cast((u >> 23) & 0xFF) - 127 + 15; + uint32_t mant = u & 0x7FFFFF; + if (((u >> 23) & 0xFF) == 0xFF) + return static_cast(sign | 0x7C00 | (mant ? 0x200 | (mant >> 13) : 0)); + if (exp >= 0x1F) return static_cast(sign | 0x7C00); + if (exp <= 0) { + if (exp < -10) return sign; + mant |= 0x800000; + uint32_t shift = static_cast(14 - exp); + uint32_t half = mant >> shift; + uint32_t rem = mant & ((1u << shift) - 1); + uint32_t mid = 1u << (shift - 1); + if (rem > mid || (rem == mid && (half & 1))) ++half; + return static_cast(sign | half); + } + uint32_t half = static_cast(exp << 10) | (mant >> 13); + uint32_t rem = mant & 0x1FFF; + if (rem > 0x1000 || (rem == 0x1000 && (half & 1))) ++half; + return static_cast(sign | half); +} +__device__ inline int DNearestInt(float fval) { + float val = fval + 12582912.0f; + int i = __float_as_int(val); + return (i & 0x007fffff) - 0x00400000; +} +__device__ inline float DLoadAct(const void* base, ActDT dt, int64_t idx) { + switch (dt) { + case ActDT::kF32: return static_cast(base)[idx]; + case ActDT::kF16: return DF16ToF32(static_cast(base)[idx]); + default: return DBF16ToF32(static_cast(base)[idx]); + } +} +__device__ __forceinline__ int GetIntB2(const int8_t* qs, int i32) { + const uint16_t* x16 = reinterpret_cast(qs); + return static_cast(x16[2 * i32 + 0]) | (static_cast(x16[2 * i32 + 1]) << 16); +} + +// Signed 8-bit x4 dot-product-accumulate, bit-identical to __dp4a (integer +// math is exact either way). The HW dot instruction (v_dot4_i32_i8 / +// __ockl_sdot4) is a perf lever, not a correctness requirement. +__device__ __forceinline__ int Dp4a(int a, int b, int acc) { + const int8_t* a8 = reinterpret_cast(&a); + const int8_t* b8 = reinterpret_cast(&b); + return acc + a8[0] * b8[0] + a8[1] * b8[1] + a8[2] * b8[2] + a8[3] * b8[3]; +} + +// ---- activation quantizers ---- +// Q8_0 (thread-per-32-block): cuda_quant_dot.cu:869. +__global__ void QuantizeQ8_0K(BlockQ8_0* __restrict__ scratch, const void* __restrict__ a, + ActDT adt, int64_t a_rs, int64_t m, int64_t nb) { + const int64_t t = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (t >= m * nb) return; + const int64_t i = t / nb; + const int64_t b = t % nb; + const int64_t elem0 = i * a_rs + b * kQK8_0; + float amax = 0.0f; + for (int j = 0; j < kQK8_0; ++j) { + const float av = fabsf(DLoadAct(a, adt, elem0 + j)); + amax = amax > av ? amax : av; + } + BlockQ8_0& y = scratch[t]; + const float d = amax / 127.0f; + const float id = d != 0.0f ? 1.0f / d : 0.0f; + y.d = DF32ToF16(d); + for (int j = 0; j < kQK8_0; ++j) { + const float x0 = DLoadAct(a, adt, elem0 + j) * id; + y.qs[j] = static_cast(roundf(x0)); + } +} + +// Q8_K (thread-per-256-superblock): cuda_quant_dot.cu QuantizeQ8KKernel. +__global__ void QuantizeQ8KK(BlockQ8_K* __restrict__ scratch, const void* __restrict__ a, + ActDT adt, int64_t a_rs, int64_t m, int64_t nsb) { + const int64_t t = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (t >= m * nsb) return; + const int64_t i = t / nsb; + const int64_t sb = t % nsb; + const int64_t elem0 = i * a_rs + sb * kQK_K; + float mx = 0.0f, amax = 0.0f; + for (int j = 0; j < kQK_K; ++j) { + const float ax = fabsf(DLoadAct(a, adt, elem0 + j)); + if (ax > amax) { amax = ax; mx = DLoadAct(a, adt, elem0 + j); } + } + BlockQ8_K& y = scratch[t]; + if (amax == 0.0f) { + y.d = 0.0f; + for (int j = 0; j < kQK_K; ++j) y.qs[j] = 0; + for (int g = 0; g < kQK_K / 16; ++g) y.bsums[g] = 0; + return; + } + const float iscale = -127.0f / mx; + for (int j = 0; j < kQK_K; ++j) { + const int v = DNearestInt(iscale * DLoadAct(a, adt, elem0 + j)); + y.qs[j] = static_cast(v < 127 ? v : 127); + } + for (int g = 0; g < kQK_K / 16; ++g) { + int sum = 0; + for (int ii = 0; ii < 16; ++ii) sum += y.qs[g * 16 + ii]; + y.bsums[g] = static_cast(sum); + } + y.d = 1.0f / iscale; +} + +// ---- dot superblocks (1:1 ports) ---- +// Q8_0 x Q8_0: cuda_quant_dot.cu QuantDotGemmQ8_0 — dp4a int core. +__device__ inline float DotQ8_0(const BlockQ8_0* wb, const BlockQ8_0* ab) { + int sumi = 0; +#pragma unroll + for (int k = 0; k < kQK8_0 / 4; ++k) + sumi = Dp4a(GetIntB2(wb->qs, k), GetIntB2(ab->qs, k), sumi); + return sumi * (DF16ToF32(wb->d) * DF16ToF32(ab->d)); +} + +// Q4_K x Q8_K: cuda_quant_dot.cu DotQ4K. dp4a-vectorized, one scale per 32. +__device__ inline float DotQ4K(const BlockQ4_K* xb, const BlockQ8_K* yb) { + const uint32_t kmask1 = 0x3f3f3f3f, kmask2 = 0x0f0f0f0f, kmask3 = 0x03030303; + const uint8_t* q4 = xb->qs; + const int8_t* q8 = yb->qs; + uint32_t utmp[4]; + memcpy(utmp, xb->scales, 12); + utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4); + const uint32_t uaux = utmp[1] & kmask1; + utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4); + utmp[2] = uaux; + utmp[0] &= kmask1; + const uint8_t* scales = reinterpret_cast(&utmp[0]); + const uint8_t* mins = reinterpret_cast(&utmp[2]); + int sumi = 0; + for (int j = 0; j < kQK_K / 16; ++j) sumi += yb->bsums[j] * mins[j / 2]; + int isum = 0; + for (int sb = 0; sb < kQK_K / 32; ++sb) { + const int scale = scales[sb]; + const uint8_t* q4b = q4 + (sb / 2) * 32; + const int8_t* q8b = q8 + sb * 32; + const int shift = (sb & 1) ? 4 : 0; + int sub = 0; + for (int l = 0; l < 32; l += 4) { + const int v = (*reinterpret_cast(q4b + l) >> shift) & 0x0F0F0F0F; + sub = Dp4a(v, *reinterpret_cast(q8b + l), sub); + } + isum += scale * sub; + } + const float d = DF16ToF32(xb->d) * yb->d; + const float dmin = DF16ToF32(xb->dmin) * yb->d; + return d * isum - dmin * sumi; +} + +// Q5_K x Q8_K: cuda_quant_dot.cu DotQ5K. Q4_K nibble + a high bit from qh. +__device__ inline float DotQ5K(const BlockQ5_K* xb, const BlockQ8_K* yb) { + const uint32_t kmask1 = 0x3f3f3f3f, kmask2 = 0x0f0f0f0f, kmask3 = 0x03030303; + const uint8_t* q4 = xb->qs; + const uint8_t* hm = xb->qh; + const int8_t* q8 = yb->qs; + uint32_t utmp[4]; + memcpy(utmp, xb->scales, 12); + utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4); + const uint32_t uaux = utmp[1] & kmask1; + utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4); + utmp[2] = uaux; + utmp[0] &= kmask1; + const uint8_t* scales = reinterpret_cast(&utmp[0]); + const uint8_t* mins = reinterpret_cast(&utmp[2]); + int sumi = 0; + for (int j = 0; j < kQK_K / 16; ++j) sumi += yb->bsums[j] * mins[j / 2]; + int isum = 0; + for (int sb = 0; sb < kQK_K / 32; ++sb) { + const int scale = scales[sb]; + const uint8_t* q4b = q4 + (sb / 2) * 32; + const int8_t* q8b = q8 + sb * 32; + const int shift = (sb & 1) ? 4 : 0; + int sub = 0; + for (int l = 0; l < 32; l += 4) { + const int lo = (*reinterpret_cast(q4b + l) >> shift) & 0x0F0F0F0F; + const int hi = ((*reinterpret_cast(hm + l) >> sb) & 0x01010101) << 4; + sub = Dp4a(lo | hi, *reinterpret_cast(q8b + l), sub); + } + isum += scale * sub; + } + const float d = DF16ToF32(xb->d) * yb->d; + const float dmin = DF16ToF32(xb->dmin) * yb->d; + return d * isum - dmin * sumi; +} + +// Q6_K x Q8_K: cuda_quant_dot.cu DotQ6K. Rebuild the 6-bit quants then scalar-MAC. +__device__ inline float DotQ6K(const BlockQ6_K* xb, const BlockQ8_K* yb) { + const uint8_t* q4 = xb->ql; + const uint8_t* qh = xb->qh; + const int8_t* q8 = yb->qs; + int8_t aux8[kQK_K]; + int8_t* a = aux8; + for (int j = 0; j < kQK_K; j += 128) { + for (int l = 0; l < 32; ++l) { + a[l + 0] = static_cast(static_cast((q4[l + 0] & 0xF) | (((qh[l] >> 0) & 3) << 4)) - 32); + a[l + 32] = static_cast(static_cast((q4[l + 32] & 0xF) | (((qh[l] >> 2) & 3) << 4)) - 32); + a[l + 64] = static_cast(static_cast((q4[l + 0] >> 4) | (((qh[l] >> 4) & 3) << 4)) - 32); + a[l + 96] = static_cast(static_cast((q4[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) - 32); + } + a += 128; q4 += 64; qh += 32; + } + a = aux8; + const int8_t* q8p = q8; + int is = 0; + int32_t aux32[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + for (int j = 0; j < kQK_K / 16; ++j) { + const int scale = xb->scales[is++]; + for (int l = 0; l < 8; ++l) aux32[l] += scale * (q8p[l] * a[l]); + q8p += 8; a += 8; + for (int l = 0; l < 8; ++l) aux32[l] += scale * (q8p[l] * a[l]); + q8p += 8; a += 8; + } + const float d = DF16ToF32(xb->d) * yb->d; + int isum = 0; + for (int l = 0; l < 8; ++l) isum += aux32[l]; + return d * isum; +} + +// ---- grouped kernels ---- +template +__global__ void GroupedQ8_0K(OutT* __restrict__ out, const uint8_t* __restrict__ weight, + const BlockQ8_0* __restrict__ act, + const int32_t* __restrict__ expert_ids, int64_t P, int64_t n, + int64_t nb, size_t w_row_bytes, bool bcast) { + const int64_t warp = static_cast(blockIdx.x) * blockDim.y + threadIdx.y; + if (warp >= P * n) return; + const int64_t p = warp / n; + const int64_t j = warp % n; + const int lane = threadIdx.x; + const int64_t e = expert_ids[p]; + const uint8_t* w_row = weight + static_cast(e * n + j) * w_row_bytes; + const BlockQ8_0* a_row = act + (bcast ? 0 : p) * nb; + float partial = 0.0f; + for (int64_t b = lane; b < nb; b += 32) { + const BlockQ8_0* wb = reinterpret_cast(w_row + static_cast(b) * + sizeof(BlockQ8_0)); + partial += DotQ8_0(wb, a_row + b); + } +#pragma unroll + for (int off = 16; off > 0; off >>= 1) partial += __shfl_down_sync(0xffffffffULL, partial, off); + if (lane == 0) { + if constexpr (sizeof(OutT) == 4) out[p * n + j] = partial; + else out[p * n + j] = DF32ToBF16(partial); + } +} + +// K-quant grouped kernel body (shared by Q4_K/Q5_K/Q6_K instantiations). +// Fmt: 0=Q4_K, 1=Q5_K, 2=Q6_K. +template +__global__ void GroupedKQ8K(OutT* __restrict__ out, const uint8_t* __restrict__ weight, + const BlockQ8_K* __restrict__ act, + const int32_t* __restrict__ expert_ids, int64_t P, int64_t n, + int64_t nsb, size_t w_row_bytes, size_t w_block_bytes, bool bcast) { + const int64_t warp = static_cast(blockIdx.x) * blockDim.y + threadIdx.y; + if (warp >= P * n) return; + const int64_t p = warp / n; + const int64_t j = warp % n; + const int lane = threadIdx.x; + const int64_t e = expert_ids[p]; + const uint8_t* w_row = weight + static_cast(e * n + j) * w_row_bytes; + const BlockQ8_K* a_row = act + (bcast ? 0 : p) * nsb; + float partial = 0.0f; + for (int64_t sb = lane; sb < nsb; sb += 32) { + const void* w_sb = w_row + static_cast(sb) * w_block_bytes; + const BlockQ8_K* a_sb = a_row + sb; + if constexpr (Fmt == 2) partial += DotQ6K(static_cast(w_sb), a_sb); + else if constexpr (Fmt == 1) partial += DotQ5K(static_cast(w_sb), a_sb); + else partial += DotQ4K(static_cast(w_sb), a_sb); + } +#pragma unroll + for (int off = 16; off > 0; off >>= 1) partial += __shfl_down_sync(0xffffffffULL, partial, off); + if (lane == 0) { + if constexpr (sizeof(OutT) == 4) out[p * n + j] = partial; + else out[p * n + j] = DF32ToBF16(partial); + } +} + + +// --- non-grouped keep-quant GEMM (cuda_quant_dot.cu QuantDotGemmKernel :706) --- +// warp per (i,j); no expert indirection (w_row = weight + j*w_row_bytes). +template +__global__ void KQuantGemmK(OutT* __restrict__ out, const uint8_t* __restrict__ weight, + const BlockQ8_K* __restrict__ act, int64_t m, int64_t n, + int64_t nsb, size_t w_row_bytes, size_t w_block_bytes) { + const int64_t warp = static_cast(blockIdx.x) * blockDim.y + threadIdx.y; + if (warp >= m * n) return; + const int64_t i = warp / n; + const int64_t j = warp % n; + const int lane = threadIdx.x; + const uint8_t* w_row = weight + static_cast(j) * w_row_bytes; + const BlockQ8_K* a_row = act + i * nsb; + float partial = 0.0f; + for (int64_t sb = lane; sb < nsb; sb += 32) { + const void* w_sb = w_row + static_cast(sb) * w_block_bytes; + if constexpr (Fmt == 2) partial += DotQ6K(static_cast(w_sb), a_row + sb); + else if constexpr (Fmt == 1) partial += DotQ5K(static_cast(w_sb), a_row + sb); + else partial += DotQ4K(static_cast(w_sb), a_row + sb); + } +#pragma unroll + for (int off = 16; off > 0; off >>= 1) partial += __shfl_down_sync(0xffffffffULL, partial, off); + if (lane == 0) { + if constexpr (sizeof(OutT) == 4) out[i * n + j] = partial; + else out[i * n + j] = DF32ToBF16(partial); + } +} + +template +__global__ void Q8_0GemmK(OutT* __restrict__ out, const uint8_t* __restrict__ weight, + const BlockQ8_0* __restrict__ act, int64_t m, int64_t n, + int64_t nb) { + const int64_t warp = static_cast(blockIdx.x) * blockDim.y + threadIdx.y; + if (warp >= m * n) return; + const int64_t i = warp / n; + const int64_t j = warp % n; + const int lane = threadIdx.x; + const uint8_t* w_row = weight + static_cast(j * nb) * sizeof(BlockQ8_0); + const BlockQ8_0* a_row = act + i * nb; + float partial = 0.0f; + for (int64_t bb = lane; bb < nb; bb += 32) { + const BlockQ8_0* wb = reinterpret_cast(w_row + static_cast(bb) * + sizeof(BlockQ8_0)); + partial += DotQ8_0(wb, a_row + bb); + } +#pragma unroll + for (int off = 16; off > 0; off >>= 1) partial += __shfl_down_sync(0xffffffffULL, partial, off); + if (lane == 0) { + if constexpr (sizeof(OutT) == 4) out[i * n + j] = partial; + else out[i * n + j] = DF32ToBF16(partial); + } +} + +inline void Check(hipError_t err, const char* what) { + if (err != hipSuccess) + throw std::runtime_error(std::string("vt rocm grouped_gemm: ") + what + ": " + + hipGetErrorString(err)); +} + +// Grow-only, per-stream activation-quant scratch pool — the mirror of the +// donor's EnsureScratch + RetireGraphScratch discipline +// (src/vt/cuda/cuda_quant_dot.cu:1507, src/vt/cuda/graph_safe_scratch.h): +// hipMalloc/hipFree/hipStreamSynchronize per call are ILLEGAL under hipGraph +// stream capture (the #473/#332 decode-graph lane), and the old free-after- +// launch also needed the sync that serialized every call. The pool allocates +// stream-ordered (hipMallocAsync), never frees during the process (a captured +// graph may have baked the pointer), and needs NO synchronization: reuse is +// stream-ordered and retirement keeps every baked pointer valid. Bounded: the +// buffer grows O(log(max/min)) times over a process. +struct StreamScratch { + void* buf = nullptr; + size_t bytes = 0; +}; + +StreamScratch& ScratchFor(hipStream_t s) { + static std::mutex mu; + static std::unordered_map pools; + std::lock_guard lk(mu); + return pools[s]; +} + +void* EnsureQuantScratch(size_t need, hipStream_t s) { + StreamScratch& sc = ScratchFor(s); + if (need > sc.bytes) { + // Retire, never free (see the header note above). + Check(hipMallocAsync(&sc.buf, need, s), "quant scratch grow"); + sc.bytes = need; + } + return sc.buf; +} + +} // namespace + + +void MatmulBTQuantKernelRocm(Queue& q, Tensor& out, const Tensor& a, const Tensor& b) { + EnsureQueueDevice(q); + const int64_t m = a.shape[0], k = a.shape[1], n = b.shape[0]; + if (m == 0 || n == 0) return; + hipStream_t s = static_cast(q.handle); + constexpr int kWarpsPerBlock = 4; + dim3 block(32, kWarpsPerBlock); + const uint8_t* w = static_cast(b.data); + if (b.dtype == DType::kQ8_0) { + if (k % kQK8_0 != 0) throw std::runtime_error("vt rocm: matmul_bt_quant Q8_0: K%32!=0"); + const int64_t nb = k / kQK8_0; + BlockQ8_0* qact = static_cast(EnsureQuantScratch( + static_cast(m) * nb * sizeof(BlockQ8_0), s)); + QuantizeQ8_0K<<((m * nb + 127) / 128), 128, 0, s>>>( + qact, a.data, ActDtOf(a.dtype), a.stride[0], m, nb); + Check(hipGetLastError(), "q8_0 quant"); + const int64_t grid = (m * n + kWarpsPerBlock - 1) / kWarpsPerBlock; + if (out.dtype == DType::kF32) + Q8_0GemmK<<(grid), block, 0, s>>>(static_cast(out.data), w, qact, m, n, nb); + else + Q8_0GemmK<<(grid), block, 0, s>>>(static_cast(out.data), w, qact, m, n, nb); + Check(hipGetLastError(), "q8_0 gemm"); + return; + } + if (b.dtype == DType::kQ4_K || b.dtype == DType::kQ5_K || b.dtype == DType::kQ6_K) { + if (k % kQK_K != 0) throw std::runtime_error("vt rocm: matmul_bt_quant K-quant: K%256!=0"); + const int64_t nsb = k / kQK_K; + const size_t w_block_bytes = b.dtype == DType::kQ6_K ? sizeof(BlockQ6_K) + : b.dtype == DType::kQ5_K ? sizeof(BlockQ5_K) + : sizeof(BlockQ4_K); + const size_t w_row_bytes = static_cast(nsb) * w_block_bytes; + BlockQ8_K* qact = static_cast(EnsureQuantScratch( + static_cast(m) * nsb * sizeof(BlockQ8_K), s)); + QuantizeQ8KK<<((m * nsb + 127) / 128), 128, 0, s>>>( + qact, a.data, ActDtOf(a.dtype), a.stride[0], m, nsb); + Check(hipGetLastError(), "q8_K quant"); + const int64_t grid = (m * n + kWarpsPerBlock - 1) / kWarpsPerBlock; + const int fmt = b.dtype == DType::kQ6_K ? 2 : b.dtype == DType::kQ5_K ? 1 : 0; + auto launch = [&](auto ot) { + using OutT = decltype(ot); + auto* o = static_cast(out.data); + if (fmt == 2) KQuantGemmK<<(grid), block, 0, s>>>(o, w, qact, m, n, nsb, w_row_bytes, w_block_bytes); + else if (fmt == 1) KQuantGemmK<<(grid), block, 0, s>>>(o, w, qact, m, n, nsb, w_row_bytes, w_block_bytes); + else KQuantGemmK<<(grid), block, 0, s>>>(o, w, qact, m, n, nsb, w_row_bytes, w_block_bytes); + }; + if (out.dtype == DType::kF32) launch(float{}); else launch(uint16_t{}); + Check(hipGetLastError(), "K-quant gemm"); + return; + } + throw std::runtime_error("vt rocm: matmul_bt_quant: unsupported weight dtype (ported: Q8_0/Q4_K/Q5_K/Q6_K; owed: Q4_0/Q2_K/Q3_K/IQ2_XXS/IQ3_XXS/IQ2_S/MXFP4 -- the loader pre-filters to the ported set, so reaching here is a bug)"); +} + + +// kMatmulBTQuantGrouped for ROCm: Q8_0 / Q4_K / Q6_K natively (the formats the +// target GDN-MoE GGUFs carry); anything else throws loudly (never a silent +// CPU-pointer deref on a discrete card). +void MatmulBTQuantGroupedKernelRocm(Queue& q, Tensor& out, const Tensor& act, + const Tensor& weight, const Tensor& expert_ids) { + EnsureQueueDevice(q); + const int64_t P = out.shape[0], n = out.shape[1], k = act.shape[1]; + if (P == 0 || n == 0) return; + const int64_t Pa = act.shape[0]; + const bool bcast = (Pa == 1 && P > 1); + hipStream_t s = static_cast(q.handle); + const uint8_t* w = static_cast(weight.data); + const int32_t* eids = static_cast(expert_ids.data); + constexpr int kWarpsPerBlock = 4; + dim3 block(32, kWarpsPerBlock); + + if (weight.dtype == DType::kQ8_0) { + if (k % kQK8_0 != 0) + throw std::runtime_error("vt rocm: matmul_bt_quant_grouped Q8_0: K must be a multiple of 32"); + const int64_t nb = k / kQK8_0; + const size_t w_row_bytes = static_cast(nb) * sizeof(BlockQ8_0); + BlockQ8_0* qact = static_cast(EnsureQuantScratch( + static_cast(Pa) * nb * sizeof(BlockQ8_0), s)); + constexpr int kQBlock = 128; + QuantizeQ8_0K<<((Pa * nb + kQBlock - 1) / kQBlock), kQBlock, 0, s>>>( + qact, act.data, ActDtOf(act.dtype), act.stride[0], Pa, nb); + Check(hipGetLastError(), "q8_0 quant"); + const int64_t grid = (P * n + kWarpsPerBlock - 1) / kWarpsPerBlock; + if (out.dtype == DType::kF32) + GroupedQ8_0K<<(grid), block, 0, s>>>( + static_cast(out.data), w, qact, eids, P, n, nb, w_row_bytes, bcast); + else + GroupedQ8_0K<<(grid), block, 0, s>>>( + static_cast(out.data), w, qact, eids, P, n, nb, w_row_bytes, bcast); + Check(hipGetLastError(), "q8_0 grouped"); + return; + } + + if (weight.dtype == DType::kQ4_K || weight.dtype == DType::kQ5_K || weight.dtype == DType::kQ6_K) { + if (k % kQK_K != 0) + throw std::runtime_error("vt rocm: matmul_bt_quant_grouped K-quant: K must be a multiple of 256"); + const int64_t nsb = k / kQK_K; + const size_t w_block_bytes = weight.dtype == DType::kQ4_K ? sizeof(BlockQ4_K) + : weight.dtype == DType::kQ5_K ? sizeof(BlockQ5_K) + : sizeof(BlockQ6_K); + const size_t w_row_bytes = static_cast(nsb) * w_block_bytes; + BlockQ8_K* qact = static_cast(EnsureQuantScratch( + static_cast(Pa) * nsb * sizeof(BlockQ8_K), s)); + QuantizeQ8KK<<((Pa * nsb + 127) / 128), 128, 0, s>>>( + qact, act.data, ActDtOf(act.dtype), act.stride[0], Pa, nsb); + Check(hipGetLastError(), "q8_K quant"); + const int64_t grid = (P * n + kWarpsPerBlock - 1) / kWarpsPerBlock; + const int fmt = weight.dtype == DType::kQ6_K ? 2 : weight.dtype == DType::kQ5_K ? 1 : 0; + auto launch = [&](auto ot) { + using OutT = decltype(ot); + auto* o = static_cast(out.data); + if (fmt == 2) GroupedKQ8K<<(grid), block, 0, s>>>(o, w, qact, eids, P, n, nsb, w_row_bytes, w_block_bytes, bcast); + else if (fmt == 1) GroupedKQ8K<<(grid), block, 0, s>>>(o, w, qact, eids, P, n, nsb, w_row_bytes, w_block_bytes, bcast); + else GroupedKQ8K<<(grid), block, 0, s>>>(o, w, qact, eids, P, n, nsb, w_row_bytes, w_block_bytes, bcast); + }; + if (out.dtype == DType::kF32) launch(float{}); else launch(uint16_t{}); + Check(hipGetLastError(), "K-quant grouped"); + return; + } + + throw std::runtime_error( + "vt rocm: matmul_bt_quant_grouped: unsupported weight dtype (ported: Q8_0/Q4_K/Q5_K/Q6_K; " + "owed: Q4_0/Q2_K/Q3_K/IQ2_XXS/IQ3_XXS/IQ2_S/MXFP4 -- the loader pre-filters, so reaching here is a bug)"); +} + +} // namespace vt::rocm diff --git a/src/vt/rocm/rocm_moe_chain.hip b/src/vt/rocm/rocm_moe_chain.hip new file mode 100644 index 000000000..c33424e7c --- /dev/null +++ b/src/vt/rocm/rocm_moe_chain.hip @@ -0,0 +1,206 @@ +// ROCm MoE combine/gate ops (BACKEND-ROCM; the model-path MoE chain, #41). +// Hand-translations from src/vt/cuda/cuda_moe.cu (MoeCombineKernel :473, +// MoeCombineGateKernel :555) and the SharedExpertGate CPU oracle +// (cpu_ops.cpp:2387), readable side by side against the donors. All f32 math; +// bf16/f32 dtype arms via the Ld/St boundary conversions. +// +// SharedExpertGate: out[t,c] = sigmoid(gl[t]) * sd[t,c] +// MoeCombine: out[r,c] = sum_j w[r,j]*expert_out[(r*k+j),c] (+ shared[r,c]) +// MoeCombineGate: MoeCombine + the shared-expert sigmoid gate folded in +// (rounded through bf16 exactly as the donor). + +#include +#include + +#include +#include +#include + +#include "vt/ops.h" + +namespace vt::rocm { +namespace { + +constexpr int kBlock = 256; + +inline void Check(hipError_t err, const char* what) { + if (err != hipSuccess) + throw std::runtime_error(std::string("vt rocm moe_chain: ") + what + ": " + + hipGetErrorString(err)); +} +inline hipStream_t AsStream(const Queue& q) { return static_cast(q.handle); } +inline unsigned GridFor(int64_t n) { + if (n <= 0) return 1; + const int64_t g = (n + kBlock - 1) / kBlock; + return static_cast(g > 65535 ? 65535 : g); +} +__device__ inline float Ld(const float* p, int64_t i) { return p[i]; } +__device__ inline float Ld(const __hip_bfloat16* p, int64_t i) { + return __bfloat162float(p[i]); +} +__device__ inline void St(float* p, int64_t i, float v) { p[i] = v; } +__device__ inline void St(__hip_bfloat16* p, int64_t i, float v) { + p[i] = __float2bfloat16(v); +} +__device__ inline float SigmoidF(float x) { return 1.0f / (1.0f + expf(-x)); } + +// SharedExpertGate (cpu_ops.cpp:2387): out[t,c] = sigmoid(gl[t]) * sd[t,c]. +// out bf16 [T,H]; sd f32 [T,H]; gl f32 [T]. +template +__global__ void SharedExpertGateK(Tout* out, const Tsd* sd, const float* gl, int64_t t, + int64_t h) { + const int64_t n = t * h; + const int64_t step = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + idx < n; idx += step) { + const int64_t row = idx / h; + St(out, idx, SigmoidF(gl[row]) * Ld(sd, idx)); + } +} + +// MoeCombine (cuda_moe.cu:473). +template +__global__ void MoeCombineK(Tout* out, const Teo* expert_out, const float* weights, + const Tsh* shared, int64_t t, int64_t h, int k, + float routed_scale) { + const int64_t n = t * h; + const int64_t step = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; idx < n; + idx += step) { + const int64_t row = idx / h; + const int64_t col = idx % h; + float acc = 0.0f; + for (int j = 0; j < k; ++j) + acc += weights[row * k + j] * Ld(expert_out, (row * k + j) * h + col); + // routed_scale multiplies the ROUTED sum only, BEFORE the shared term is + // added (upstream apply_routed_scale_to_output: fused_output *= factor, + // shared untouched) — one standalone f32 multiply on the finished + // accumulator, bit-identical to the CPU reference under -ffp-contract=off. + acc *= routed_scale; + if (shared != nullptr) acc += Ld(shared, idx); + St(out, idx, acc); + } +} + +// MoeCombineGate (cuda_moe.cu:555): MoeCombine + shared-expert sigmoid gate +// folded in, the shared term rounded through bf16 exactly as the donor. +template +__global__ void MoeCombineGateK(Tout* out, const Teo* expert_out, const float* weights, + const Tsd* sd, const float* gl, int64_t t, int64_t h, + int k) { + const int64_t n = t * h; + const int64_t step = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; idx < n; + idx += step) { + const int64_t row = idx / h; + const int64_t col = idx % h; + float acc = 0.0f; + for (int j = 0; j < k; ++j) + acc += weights[row * k + j] * Ld(expert_out, (row * k + j) * h + col); + const float sv = SigmoidF(gl[row]) * Ld(sd, idx); + acc += __bfloat162float(__float2bfloat16(sv)); + St(out, idx, acc); + } +} + +} // namespace + +void SharedExpertGateKernelRocm(Queue& q, Tensor& out, const Tensor& sd, const Tensor& gl) { + // Same f16 hazard as the combine ops (review sweep on #509): the seam admits + // any float dtype and Tensor::Ptr() is unchecked. + VT_CHECK(sd.dtype == DType::kF32 || sd.dtype == DType::kBF16, + "rocm shared_expert_gate: unsupported sd dtype (f32/bf16 only)"); + VT_CHECK(out.dtype == DType::kF32 || out.dtype == DType::kBF16, + "rocm shared_expert_gate: unsupported out dtype (f32/bf16 only)"); + const int64_t t = out.shape[0], h = out.shape[1]; + if (t == 0 || h == 0) return; + hipStream_t s = AsStream(q); + const int64_t n = t * h; + const bool obf = out.dtype == DType::kBF16, sbf = sd.dtype == DType::kBF16; + if (obf && sbf) + SharedExpertGateK<<>>(out.Ptr<__hip_bfloat16>(), + sd.Ptr<__hip_bfloat16>(), + gl.Ptr(), t, h); + else if (obf) + SharedExpertGateK<<>>(out.Ptr<__hip_bfloat16>(), sd.Ptr(), + gl.Ptr(), t, h); + else if (sbf) + SharedExpertGateK<<>>(out.Ptr(), sd.Ptr<__hip_bfloat16>(), + gl.Ptr(), t, h); + else + SharedExpertGateK<<>>(out.Ptr(), sd.Ptr(), + gl.Ptr(), t, h); + Check(hipGetLastError(), "shared_expert_gate launch"); +} + +void MoeCombineKernelRocm(Queue& q, Tensor& out, const Tensor& expert_out, + const Tensor& weights, const Tensor* shared, float routed_scale) { + // The donor's dtype refusals (cuda_moe.cu:520-524) — f16 is admitted by the + // seam's IsFloat gate, and Tensor::Ptr() is an unchecked cast, so without + // these an f16 expert_out would be read at 4 bytes/element from a 2-byte + // allocation (review sweep on #509). + VT_CHECK(expert_out.dtype == DType::kF32 || expert_out.dtype == DType::kBF16, + "rocm moe_combine: unsupported expert_out dtype (f32/bf16 only)"); + VT_CHECK(out.dtype == DType::kF32 || out.dtype == DType::kBF16, + "rocm moe_combine: unsupported out dtype (f32/bf16 only)"); + VT_CHECK(shared == nullptr || shared->dtype == DType::kF32 || + shared->dtype == DType::kBF16, + "rocm moe_combine: unsupported shared dtype (f32/bf16 only)"); + const int64_t t = out.shape[0], h = out.shape[1]; + const int k = static_cast(weights.shape[1]); + const int64_t n = t * h; + if (n == 0) return; + hipStream_t s = AsStream(q); + auto launch = [&](auto eo, auto sh, auto ot) { + using Teo = decltype(eo); using Tsh = decltype(sh); using Tout = decltype(ot); + MoeCombineK<<>>( + out.Ptr(), expert_out.Ptr(), weights.Ptr(), + shared != nullptr ? shared->Ptr() : nullptr, t, h, k, routed_scale); + }; + auto by_shared = [&](auto eo, auto ot) { + if (shared != nullptr && shared->dtype == DType::kBF16) launch(eo, __hip_bfloat16{}, ot); + else launch(eo, float{}, ot); + }; + auto by_out = [&](auto eo) { + if (out.dtype == DType::kBF16) by_shared(eo, __hip_bfloat16{}); + else by_shared(eo, float{}); + }; + if (expert_out.dtype == DType::kBF16) by_out(__hip_bfloat16{}); + else by_out(float{}); + Check(hipGetLastError(), "moe_combine launch"); +} + +void MoeCombineGateKernelRocm(Queue& q, Tensor& out, const Tensor& expert_out, + const Tensor& weights, const Tensor& sd, const Tensor& gl) { + // Same donor refusals as MoeCombine (cuda_moe.cu:597-604). + VT_CHECK(expert_out.dtype == DType::kF32 || expert_out.dtype == DType::kBF16, + "rocm moe_combine_gate: unsupported expert_out dtype (f32/bf16 only)"); + VT_CHECK(out.dtype == DType::kF32 || out.dtype == DType::kBF16, + "rocm moe_combine_gate: unsupported out dtype (f32/bf16 only)"); + VT_CHECK(sd.dtype == DType::kF32 || sd.dtype == DType::kBF16, + "rocm moe_combine_gate: unsupported sd dtype (f32/bf16 only)"); + const int64_t t = out.shape[0], h = out.shape[1]; + const int k = static_cast(weights.shape[1]); + const int64_t n = t * h; + if (n == 0) return; + hipStream_t s = AsStream(q); + auto launch = [&](auto eo, auto sd_t, auto ot) { + using Teo = decltype(eo); using Tsd = decltype(sd_t); using Tout = decltype(ot); + MoeCombineGateK<<>>( + out.Ptr(), expert_out.Ptr(), weights.Ptr(), sd.Ptr(), + gl.Ptr(), t, h, k); + }; + auto by_sd = [&](auto eo, auto ot) { + if (sd.dtype == DType::kBF16) launch(eo, __hip_bfloat16{}, ot); + else launch(eo, float{}, ot); + }; + auto by_out = [&](auto eo) { + if (out.dtype == DType::kBF16) by_sd(eo, __hip_bfloat16{}); + else by_sd(eo, float{}); + }; + if (expert_out.dtype == DType::kBF16) by_out(__hip_bfloat16{}); + else by_out(float{}); + Check(hipGetLastError(), "moe_combine_gate launch"); +} + +} // namespace vt::rocm diff --git a/src/vt/rocm/rocm_ops.hip b/src/vt/rocm/rocm_ops.hip index aed29660c..3868b8e83 100644 --- a/src/vt/rocm/rocm_ops.hip +++ b/src/vt/rocm/rocm_ops.hip @@ -54,6 +54,17 @@ void ApplyLogitBiasKernelRocm(Queue& q, Tensor& logits, const Tensor& rows, cons const Tensor& biases); // Companion MoE-path op (same TU): elementwise silu(gate)*up. void MoeSiluMulKernelRocm(Queue& q, Tensor& out, const Tensor& gate, const Tensor& up); +// MoE-path combine/gate ops (rocm_moe_chain.hip): shared-expert gate and the +// weighted expert combinations. +void SharedExpertGateKernelRocm(Queue& q, Tensor& out, const Tensor& sd, const Tensor& gl); +void MoeCombineKernelRocm(Queue& q, Tensor& out, const Tensor& expert_out, + const Tensor& weights, const Tensor* shared, float routed_scale); +void MoeCombineGateKernelRocm(Queue& q, Tensor& out, const Tensor& expert_out, + const Tensor& weights, const Tensor& sd, const Tensor& gl); +// Grouped quant expert GEMM (rocm_grouped_gemm.hip): Q8_0/Q4_K/Q5_K/Q6_K native. +void MatmulBTQuantKernelRocm(Queue& q, Tensor& out, const Tensor& a, const Tensor& b); +void MatmulBTQuantGroupedKernelRocm(Queue& q, Tensor& out, const Tensor& act, + const Tensor& weight, const Tensor& expert_ids); // BACKEND-ROCM-GDN-KERNELS family 1 (rocm_gdn_state.hip): the indexed state I/O // pair `IndexedGdnOpsNative()` requires (issue #41, spec rocm-gdn-kernels.md). void GdnStateGatherKernelRocm(Queue& q, Tensor& working, const Tensor& cache, @@ -181,6 +192,19 @@ struct Registrar { static_cast(&ApplyTokenMaskKernelRocm))); RegisterOp(OpId::kMoeSiluMul, DeviceType::kROCM, reinterpret_cast(static_cast(&MoeSiluMulKernelRocm))); + RegisterOp(OpId::kSharedExpertGate, DeviceType::kROCM, + reinterpret_cast( + static_cast(&SharedExpertGateKernelRocm))); + RegisterOp(OpId::kMoeCombine, DeviceType::kROCM, + reinterpret_cast(static_cast(&MoeCombineKernelRocm))); + RegisterOp(OpId::kMoeCombineGate, DeviceType::kROCM, + reinterpret_cast( + static_cast(&MoeCombineGateKernelRocm))); + RegisterOp(OpId::kMatmulBTQuant, DeviceType::kROCM, + reinterpret_cast(static_cast(&MatmulBTQuantKernelRocm))); + RegisterOp(OpId::kMatmulBTQuantGrouped, DeviceType::kROCM, + reinterpret_cast( + static_cast(&MatmulBTQuantGroupedKernelRocm))); RegisterOp(OpId::kGdnStateGather, DeviceType::kROCM, reinterpret_cast( static_cast(&GdnStateGatherKernelRocm))); diff --git a/tests/vllm/test_gguf_keep_quant.cpp b/tests/vllm/test_gguf_keep_quant.cpp index e299e5f32..d5fb35ab9 100644 --- a/tests/vllm/test_gguf_keep_quant.cpp +++ b/tests/vllm/test_gguf_keep_quant.cpp @@ -59,8 +59,8 @@ using vllm::RouteGgufTensor; namespace { // ggml type ids (ggml/include/ggml.h:390-432). -constexpr uint32_t kF32 = 0, kF16 = 1, kQ4_0 = 2, kQ8_0 = 8, kQ3_K = 11, - kQ4_K = 12, kQ5_K = 13, kQ6_K = 14, kQ8_K = 15, +constexpr uint32_t kF32 = 0, kF16 = 1, kQ4_0 = 2, kQ8_0 = 8, kQ2_K = 10, + kQ3_K = 11, kQ4_K = 12, kQ5_K = 13, kQ6_K = 14, kQ8_K = 15, kIQ2_S = 22, kIQ4_XS = 23, kBF16 = 30, kMXFP4 = 39; // Every executable weight encoding, with a K that is a whole number of blocks. @@ -232,6 +232,40 @@ TEST_CASE("keep-quant expert split is lossless per expert") { } } +TEST_CASE("keep-quant routing respects the RUNNING DEVICE's format set (review #523)") { + // Registering kMatmulBTQuant flips keep-quant loader-wide via the boolean + // GgufQuantComputeAvailable(), but a device's kernel set can be narrower + // than the CPU admission list. On ROCm exactly {Q8_0, Q4_K, Q5_K, Q6_K} are + // implemented; with no CPU fallback tier on a discrete card, an unsupported + // format that flipped to keep-quant would throw at FORWARD time with the + // model fully resident. The loader must keep the pre-existing expand_bf16 + // residency for those formats instead. + const vt::DeviceType dev = + vllm::platforms::CurrentPlatform().device_type(); + if (dev != vt::DeviceType::kROCM) { + MESSAGE("non-ROCm host (the device set is full there); the device-gated " + "arms are asserted on gfx1100"); + return; + } + const std::vector shape = {4, 256}; // [out, in]: K = shape[1] = 256 elems, whole blocks + const auto route = [&](uint32_t ty) { + return RouteGgufTensor(/*keep_quant=*/true, /*keep_f16=*/true, + /*nvfp4_fp4=*/false, /*cpu_ref=*/false, + GgufTensorRole::kMatmulWeight, ty, shape); + }; + // The supported set keeps quant residency (ggml type ids per the constants + // at the top of this file). + CHECK(route(kQ4_0) == GgufResidency::kExpandBf16); // unsupported -> expand + CHECK(route(kQ8_0) == GgufResidency::kKeepQuant); + CHECK(route(kQ4_K) == GgufResidency::kKeepQuant); + CHECK(route(kQ5_K) == GgufResidency::kKeepQuant); + CHECK(route(kQ6_K) == GgufResidency::kKeepQuant); + CHECK(route(kQ2_K) == GgufResidency::kExpandBf16); // owed, not silently kept + // keep-f16 must be OFF on ROCm: MatmulBTKernelRocm accepts bf16/f32 only. + const GgufLoadPolicy pol = GgufLoadPolicy::FromEnv(); + CHECK(!pol.keep_f16); +} + TEST_CASE("keep-quant residency refuses ragged K and out-of-span slices") { const int64_t n = 2, k = 64; const size_t nbytes = BlockBytesFor(kQ8_0, n * k); @@ -305,9 +339,19 @@ TEST_CASE("routing table is TOTAL: every role x every encoding is explicit") { // --- the independent expectation --- // IQ2_S (256-elem, Q8_K-act) and MXFP4 (32-elem, Q8_0-act) are keep-quant // capable as of the UD-IQ2_M vehicle, so they route like the others. - const bool block_capable = + // The DEVICE axis (review #523): the running device's kernel set can be + // narrower than the loader's CPU-derived list — ROCm implements exactly + // {Q8_0, Q4_K, Q5_K, Q6_K}; the rest keep expand_bf16 there. + const bool cpu_capable = type == kQ4_0 || type == kQ8_0 || type == kQ3_K || type == kQ4_K || type == kQ5_K || type == kQ6_K || type == kIQ2_S || type == kMXFP4; + const bool rocm = + vllm::platforms::CurrentPlatform().device_type() == + vt::DeviceType::kROCM; + const bool device_capable = + !rocm || type == kQ8_0 || type == kQ4_K || type == kQ5_K || + type == kQ6_K; + const bool block_capable = cpu_capable && device_capable; const int64_t blk = (type == kQ4_0 || type == kQ8_0 || type == kMXFP4) ? 32 : 256; bool expect_keep = false; @@ -340,9 +384,13 @@ TEST_CASE("routing table is TOTAL: every role x every encoding is explicit") { } } // Both outcomes are actually exercised (a table that never keeps anything - // would pass every assertion above vacuously). - CHECK(kept == 16); // 8 block-capable encodings x 2 keep-capable roles - CHECK(expanded == 13 * 36 - 16); // 13 types x (6 roles x 6 shapes) - kept + // would pass every assertion above vacuously). The kept count is + // device-dependent (review #523): 8 block-capable encodings x 2 keep-capable + // roles where the device covers the CPU list; 4 x 2 on ROCm. + const bool rocm_host = + vllm::platforms::CurrentPlatform().device_type() == vt::DeviceType::kROCM; + CHECK(kept == (rocm_host ? 8 : 16)); + CHECK(expanded == 13 * 36 - (rocm_host ? 8 : 16)); } TEST_CASE("tensors that are value- or layout-rewritten NEVER keep quant") { @@ -366,6 +414,14 @@ TEST_CASE("tensors that are value- or layout-rewritten NEVER keep quant") { TEST_CASE("GgufLoadPolicy::FromEnv reads VT_CPU_REF and VT_GGUF_KEEP_QUANT") { ::unsetenv("VT_CPU_REF"); ::unsetenv("VT_GGUF_KEEP_QUANT"); + // keep_f16 additionally requires an f16-capable MatmulBT on the running + // device (review #523): the ROCm kernel accepts bf16/f32 only, so keep_f16 + // is OFF on ROCm regardless of expand_nk. + const bool f16_device_ok = + vllm::platforms::CurrentPlatform().device_type() != vt::DeviceType::kROCM; + const auto keep_f16_expected = [&](const GgufLoadPolicy& q) { + return q.expand_nk && f16_device_ok; + }; { // PRODUCTION DEFAULT SINCE CIQ G4: keep-quant follows the running device's // ability to EXECUTE the quantized GEMM. The expectation is derived from @@ -380,7 +436,7 @@ TEST_CASE("GgufLoadPolicy::FromEnv reads VT_CPU_REF and VT_GGUF_KEEP_QUANT") { // repack-source release + load-time prefault removed L6's two objections // (RSS-neutral, prefill regression), so it measures 1.01x llama.cpp RSS with // prefill/decode at-or-ahead and byte-identical tokens. - CHECK(p.keep_f16 == vllm::GgufQuantComputeAvailable()); + CHECK(p.keep_f16 == (vllm::GgufQuantComputeAvailable() && f16_device_ok)); CHECK_FALSE(p.cpu_ref); } ::setenv("VT_GGUF_KEEP_QUANT", "1", 1); @@ -390,7 +446,7 @@ TEST_CASE("GgufLoadPolicy::FromEnv reads VT_CPU_REF and VT_GGUF_KEEP_QUANT") { // NB compare to expand_nk, NOT GgufQuantComputeAvailable(): with keep-quant // env-forced, expand_nk holds even on a CUDA build where the quant GEMM is // unregistered (GgufQuantComputeAvailable() is false there). - CHECK(GgufLoadPolicy::FromEnv().keep_f16 == GgufLoadPolicy::FromEnv().expand_nk); + CHECK(GgufLoadPolicy::FromEnv().keep_f16 == keep_f16_expected(GgufLoadPolicy::FromEnv())); // The opt-out must work after the default flip. ::setenv("VT_GGUF_KEEP_F16", "0", 1); CHECK_FALSE(GgufLoadPolicy::FromEnv().keep_f16); @@ -407,11 +463,11 @@ TEST_CASE("GgufLoadPolicy::FromEnv reads VT_CPU_REF and VT_GGUF_KEEP_QUANT") { // it is inert with keep-quant off (nothing to keep) or under VT_CPU_REF. ::setenv("VT_GGUF_KEEP_QUANT", "1", 1); ::setenv("VT_GGUF_KEEP_F16", "1", 1); - CHECK(GgufLoadPolicy::FromEnv().keep_f16 == GgufLoadPolicy::FromEnv().expand_nk); + CHECK(GgufLoadPolicy::FromEnv().keep_f16 == keep_f16_expected(GgufLoadPolicy::FromEnv())); for (const char* on : {"1", "true", "on"}) { ::setenv("VT_GGUF_KEEP_F16", on, 1); CAPTURE(on); - CHECK(GgufLoadPolicy::FromEnv().keep_f16 == GgufLoadPolicy::FromEnv().expand_nk); + CHECK(GgufLoadPolicy::FromEnv().keep_f16 == keep_f16_expected(GgufLoadPolicy::FromEnv())); } ::unsetenv("VT_GGUF_KEEP_F16"); ::setenv("VT_GGUF_KEEP_QUANT", "1", 1); diff --git a/tests/vt/test_backend_cross_device.cpp b/tests/vt/test_backend_cross_device.cpp index 186500e77..5d8d9530d 100644 --- a/tests/vt/test_backend_cross_device.cpp +++ b/tests/vt/test_backend_cross_device.cpp @@ -1980,6 +1980,308 @@ TEST_CASE("MoeRouterTopK matches the CPU oracle (f32 and bf16 logits)") { } } +// Scalar bf16 RNE round-trip helpers for host-side oracles (the MoE combine +// gate reference rounds the shared term through bf16 exactly like the kernel). +static uint16_t F32ToBf16Rne(float f) { + uint32_t u; + std::memcpy(&u, &f, 4); + return static_cast((u + 0x7FFFu + ((u >> 16) & 1u)) >> 16); +} +static float Bf16ToF32(uint16_t b) { + uint32_t u = static_cast(b) << 16; + float f; + std::memcpy(&f, &u, 4); + return f; +} + +TEST_CASE("MoE combine/gate ops match the CPU oracle") { + constexpr int64_t T = 5, H = 64, K = 3; + const size_t en = static_cast(T) * K * H, on = static_cast(T) * H; + const std::vector eo = RandomVec(en, 911); + const std::vector w = RandomVec(static_cast(T) * K, 912, 0.0f, 1.0f); + const std::vector sd = RandomVec(on, 913); + const std::vector eo_bf = Bf16Bits(eo), sd_bf = Bf16Bits(sd); + // SharedExpertGate (sigmoid*mul), MoeCombine (weighted expert sum +/- + // shared), MoeCombineGate (combine + folded shared gate). f32 and bf16 arms, + // PLUS the production dtype mix the model actually runs (review sweep on + // #509): expert_out bf16 (qwen3_5.cpp DBuf ddown), shared bf16, out bf16. + // MoeCombine/MoeCombineGate are thread-per-element with a single store + // rounding and NO cross-lane reduction (cuda_moe.cu:465-468), so both arms + // are asserted BIT-EXACT — the NMSE aggregate would tolerate a few wrong + // elements, which is exactly how a 2x OOB read hides. + const std::vector gl = RandomVec(static_cast(T), 914); + + for (DeviceType dt : RegisteredDevices()) { + vt::Backend& dev = vt::GetBackend(dt); + Queue q = dev.CreateQueue(); + const Device d{dt, 0}; + // CPU oracle for all three, f32. + std::vector ref_sg_b(on, 0); + std::vector ref_c(on), ref_cg(on); + { + vt::Backend& cpu = vt::GetBackend(DeviceType::kCPU); + Queue cq = cpu.CreateQueue(); + const Device cd{DeviceType::kCPU, 0}; + std::vector csd = sd, cgl = gl, ceo = eo, cw = w; + Tensor tout = Tensor::Contiguous(ref_sg_b.data(), DType::kBF16, cd, {T, H}); + Tensor tsd = T2(csd.data(), cd, T, H); + Tensor tgl = T1(cgl.data(), cd, T); + if (OpAvailable(vt::OpId::kSharedExpertGate, DeviceType::kCPU)) + vt::SharedExpertGate(cq, tout, tsd, tgl); + Tensor teo = Tensor::Contiguous(ceo.data(), DType::kF32, cd, {T, K, H}); + Tensor tw = T2(cw.data(), cd, T, K); + Tensor to2 = T2(ref_c.data(), cd, T, H); + if (OpAvailable(vt::OpId::kMoeCombine, DeviceType::kCPU)) + vt::MoeCombine(cq, to2, teo, tw, &tsd, 1.0f); + cpu.DestroyQueue(cq); + // MoeCombineGate has no CPU op registration; the oracle is the composite + // computed on host: MoeCombine (no shared) + bf16-round(sigmoid(gl)*sd). + for (int64_t r = 0; r < T; ++r) { + const float g = 1.0f / (1.0f + std::exp(-gl[static_cast(r)])); + for (int64_t c2 = 0; c2 < H; ++c2) { + float acc = 0.0f; + for (int64_t j = 0; j < K; ++j) + acc += w[static_cast(r * K + j)] * eo[static_cast((r * K + j) * H + c2)]; + const float sv = g * sd[static_cast(r * H + c2)]; + const uint16_t svb = F32ToBf16Rne(sv); + acc += Bf16ToF32(svb); + ref_cg[static_cast(r * H + c2)] = acc; + } + } + } + // device + DevBuf deo(dev, q, en), dw(dev, q, T * K), dsd(dev, q, on), dgl(dev, q, T), dout(dev, q, on); + DevBufBytes doutb(dev, q, on * 2); + deo.Upload(eo); dw.Upload(w); dsd.Upload(sd); dgl.Upload(gl); + Tensor teo = Tensor::Contiguous(deo.ptr(), DType::kF32, d, {T, K, H}); + Tensor tw = T2(dw.ptr(), d, T, K); + Tensor tsd = T2(dsd.ptr(), d, T, H); + Tensor tgl = T1(dgl.ptr(), d, T); + Tensor tout = T2(dout.ptr(), d, T, H); + if (OpAvailable(vt::OpId::kSharedExpertGate, dt)) { + Tensor toutb = Tensor::Contiguous(doutb.ptr(), DType::kBF16, d, {T, H}); + vt::SharedExpertGate(q, toutb, tsd, tgl); + std::vector gotb(on); + doutb.Download(gotb.data()); + CHECK(gotb == ref_sg_b); // both sides store bf16: bit-exact + } + if (OpAvailable(vt::OpId::kMoeCombine, dt)) { + vt::MoeCombine(q, tout, teo, tw, &tsd, 1.0f); + // Thread-per-element, single store rounding, no cross-lane reduction: + // bit-exact is the achievable and asserted bar (review sweep on #509). + CHECK(dout.Download() == ref_c); + } + if (OpAvailable(vt::OpId::kMoeCombineGate, dt)) { + vt::MoeCombineGate(q, tout, teo, tw, tsd, tgl); + CHECK(Nmse(ref_cg, dout.Download()) <= kNmseTol); + } + + // The production bf16 arm: expert_out bf16 + shared bf16 + out bf16 + // (qwen3_5.cpp:5463 ddown / :5326 shared). The CPU oracle runs the same + // ops on the same bf16 tensors; both sides thread-per-element with the + // same sequential K order, so the assertion is BIT-EXACT. + DevBufBytes deo_bf(dev, q, en * 2), dsd_bf(dev, q, on * 2), dout_bf(dev, q, on * 2); + deo_bf.Upload(eo_bf.data()); + dsd_bf.Upload(sd_bf.data()); + Tensor teo_b = Tensor::Contiguous(deo_bf.ptr(), DType::kBF16, d, {T, K, H}); + Tensor tsd_b = Tensor::Contiguous(dsd_bf.ptr(), DType::kBF16, d, {T, H}); + Tensor tout_b = Tensor::Contiguous(dout_bf.ptr(), DType::kBF16, d, {T, H}); + if (OpAvailable(vt::OpId::kMoeCombine, dt) && + OpAvailable(vt::OpId::kMoeCombine, DeviceType::kCPU)) { + // CPU reference on bf16 tensors. + std::vector ref_b(on, 0); + { + vt::Backend& cpu = vt::GetBackend(DeviceType::kCPU); + Queue cq = cpu.CreateQueue(); + const Device cd{DeviceType::kCPU, 0}; + std::vector ceo = eo_bf, csd = sd_bf; + std::vector cw = w; + Tensor r = Tensor::Contiguous(ref_b.data(), DType::kBF16, cd, {T, H}); + Tensor teo_c = Tensor::Contiguous(ceo.data(), DType::kBF16, cd, {T, K, H}); + Tensor tw_c = T2(cw.data(), cd, T, K); + Tensor tsd_c = Tensor::Contiguous(csd.data(), DType::kBF16, cd, {T, H}); + vt::MoeCombine(cq, r, teo_c, tw_c, &tsd_c, 0.7f); + cpu.DestroyQueue(cq); + } + vt::MoeCombine(q, tout_b, teo_b, tw, &tsd_b, 0.7f); + std::vector got_b(on); + dout_bf.Download(got_b.data()); + CHECK(got_b == ref_b); + } + dev.DestroyQueue(q); + } +} + +TEST_CASE("non-grouped keep-quant GEMM (Q8_0/Q4_K/Q5_K/Q6_K) matches the CPU oracle") { + // kMatmulBTQuant (op 74) on ROCm vs the CPU keep-quant reference. The + // non-grouped arm carries PR #523's headline mechanism and had NO coverage + // (review sweep 2026-08-13); the ROCm dispatcher's src-vs-out dtype mix-up + // in the fused preamble (the 0.8B divergence, row/ROCM-GDN-08B-FIX) is + // exactly the class an untested-but-registered op hides. REQUIRE (not skip) + // on ROCm so a dropped RegisterOp can never pass silently. + constexpr int64_t M = 3, N = 8, K = 512; + struct Fmt { vt::DType dt; int64_t block_bytes; int d_off; int dmin_off; const char* name; }; + const Fmt fmts[] = { + {vt::DType::kQ8_0, 34, 0, -1, "q8_0"}, + {vt::DType::kQ4_K, 144, 0, 2, "q4_K"}, + {vt::DType::kQ6_K, 210, 208, -1, "q6_K"}, + {vt::DType::kQ5_K, 176, 0, 2, "q5_K"}, + }; + const bool rocm_present = OpAvailable(vt::OpId::kMatmulBTQuant, DeviceType::kROCM); + const bool any_rocm = [&] { + for (DeviceType dt : RegisteredDevices()) if (dt == DeviceType::kROCM) return true; + return false; + }(); + if (any_rocm) { + REQUIRE_MESSAGE(rocm_present, + "kMatmulBTQuant must be registered on ROCm (the keep-quant " + "loader flips on it) — a missing registration is a failure, " + "never a skip"); + } + for (const Fmt& f : fmts) { + CAPTURE(f.name); + const int64_t elems_per_block = (f.dt == vt::DType::kQ8_0) ? 32 : 256; + const int64_t blocks_per_row = K / elems_per_block; + const size_t row_bytes = static_cast(blocks_per_row) * f.block_bytes; + const size_t wn = static_cast(N) * row_bytes; + std::mt19937 rng(779); + std::vector wt(wn); + for (uint8_t& b : wt) b = static_cast(rng() & 0xFF); + for (int64_t r = 0; r < N; ++r) + for (int64_t bIdx = 0; bIdx < blocks_per_row; ++bIdx) { + uint8_t* blk = wt.data() + r * row_bytes + bIdx * f.block_bytes; + const float jitter = 1.0f + 0.05f * static_cast((r + bIdx) % 7); + auto put16 = [&](int off, float v) { uint16_t h = vt::F32ToF16(v); std::memcpy(blk + off, &h, 2); }; + if (f.d_off >= 0) put16(f.d_off, 0.0125f * jitter); + if (f.dmin_off >= 0) put16(f.dmin_off, 0.0075f * jitter); + } + const size_t an = static_cast(M) * K, on = static_cast(M) * N; + const std::vector act = RandomVec(an, 780, -0.5f, 0.5f); + std::vector ref(on, 0.0f); + { + vt::Backend& cpu = vt::GetBackend(DeviceType::kCPU); + Queue cq = cpu.CreateQueue(); + const Device cd{DeviceType::kCPU, 0}; + std::vector ca = act; + std::vector cw = wt; + Tensor tout = T2(ref.data(), cd, M, N); + Tensor tact = T2(ca.data(), cd, M, K); + Tensor twt = Tensor::Contiguous(cw.data(), f.dt, cd, {N, K}); + vt::MatmulBTQuant(cq, tout, tact, twt); + cpu.DestroyQueue(cq); + } + for (DeviceType dt : RegisteredDevices()) { + if (!OpAvailable(vt::OpId::kMatmulBTQuant, dt)) continue; + CAPTURE(DeviceName(dt)); + vt::Backend& dev = vt::GetBackend(dt); + Queue q = dev.CreateQueue(); + const Device d{dt, 0}; + DevBuf da(dev, q, an); + DevBufBytes dwt(dev, q, wn); + DevBuf dout(dev, q, on); + da.Upload(act); + dwt.Upload(wt.data()); + Tensor tact = T2(da.ptr(), d, M, K); + Tensor twt = Tensor::Contiguous(dwt.ptr(), f.dt, d, {N, K}); + Tensor tout = T2(dout.ptr(), d, M, N); + vt::MatmulBTQuant(q, tout, tact, twt); + CHECK(Nmse(ref, dout.Download()) <= kNmseTol); + dev.DestroyQueue(q); + } + } +} + +TEST_CASE("grouped quant expert GEMM (Q8_0/Q4_K/Q6_K) matches the CPU oracle") { + // kMatmulBTQuantGrouped on ROCm vs the CPU keep-quant reference + // (cpu_quant_gemm.cpp:305). Valid random blocks (valid f16 deltas, random + // quants) at a real expert-MLP shape. Integer cores are bit-exact ports; + // the f16/f32 scale sum reassociates across lanes, so NMSE <= 5e-4. + constexpr int64_t P = 3, N = 8, K = 512; // K%256==0 (K-quant superblocks) + constexpr int64_t E = 4; // experts + const std::vector eids = {2, 0, 3}; // routed experts (non-sorted) + + struct Fmt { vt::DType dt; int64_t block_bytes; int d_off; int dmin_off; const char* name; }; + // offsets from ggml-common.h (restated in cpu_quant_blocks.h) + const Fmt fmts[] = { + {vt::DType::kQ8_0, 34, 0, -1, "q8_0"}, // {d; qs[32]} K blocks of 32 + {vt::DType::kQ4_K, 144, 0, 2, "q4_K"}, // {d,dmin,sc,qs} superblocks of 256 + {vt::DType::kQ6_K, 210, 208, -1, "q6_K"},// {ql,qh,scales,d} superblocks of 256 + {vt::DType::kQ5_K, 176, 0, 2, "q5_K"}, // {d,dmin,sc,qh,qs} superblocks of 256 + }; + + // REQUIRE-proven registration on ROCm (never a silent skip — review sweep + // on #523: an OpAvailable-guarded case passes green with the registration + // deleted). + const bool any_rocm = [&] { + for (DeviceType dt : RegisteredDevices()) if (dt == DeviceType::kROCM) return true; + return false; + }(); + if (any_rocm) { + REQUIRE_MESSAGE(OpAvailable(vt::OpId::kMatmulBTQuantGrouped, DeviceType::kROCM), + "kMatmulBTQuantGrouped must be registered on ROCm — a missing " + "registration is a failure, never a skip"); + } + for (const Fmt& f : fmts) { + CAPTURE(f.name); + const int64_t elems_per_block = (f.dt == vt::DType::kQ8_0) ? 32 : 256; + const int64_t blocks_per_row = K / elems_per_block; + const size_t row_bytes = static_cast(blocks_per_row) * f.block_bytes; + const size_t wn = static_cast(E) * N * row_bytes; + // Build valid random blocks: random quant bytes, small positive f16 deltas. + std::mt19937 rng(777); + std::vector wt(wn); + for (uint8_t& b : wt) b = static_cast(rng() & 0xFF); + for (int64_t r = 0; r < E * N; ++r) + for (int64_t bIdx = 0; bIdx < blocks_per_row; ++bIdx) { + uint8_t* blk = wt.data() + r * row_bytes + bIdx * f.block_bytes; + const float jitter = 1.0f + 0.05f * static_cast((r + bIdx) % 7); + auto put16 = [&](int off, float v) { uint16_t h = vt::F32ToF16(v); std::memcpy(blk + off, &h, 2); }; + if (f.d_off >= 0) put16(f.d_off, 0.0125f * jitter); + if (f.dmin_off >= 0) put16(f.dmin_off, 0.0075f * jitter); + } + const size_t an = static_cast(P) * K, on = static_cast(P) * N; + const std::vector act = RandomVec(an, 778, -0.5f, 0.5f); + + std::vector ref(on, 0.0f); + { + vt::Backend& cpu = vt::GetBackend(DeviceType::kCPU); + Queue cq = cpu.CreateQueue(); + const Device cd{DeviceType::kCPU, 0}; + std::vector ca = act; + std::vector cw = wt; + std::vector ce = eids; + Tensor tout = T2(ref.data(), cd, P, N); + Tensor tact = T2(ca.data(), cd, P, K); + Tensor twt = Tensor::Contiguous(cw.data(), f.dt, cd, {E * N, K}); + Tensor te = TI32(ce.data(), cd, P); + vt::MatmulBTQuantGrouped(cq, tout, tact, twt, te); + cpu.DestroyQueue(cq); + } + for (DeviceType dt : RegisteredDevices()) { + if (!OpAvailable(vt::OpId::kMatmulBTQuantGrouped, dt)) continue; + CAPTURE(DeviceName(dt)); + vt::Backend& dev = vt::GetBackend(dt); + Queue q = dev.CreateQueue(); + const Device d{dt, 0}; + DevBuf da(dev, q, an); + DevBufBytes dwt(dev, q, wn); + DevBufI32 de(dev, q, P); + DevBuf dout(dev, q, on); + da.Upload(act); + dwt.Upload(wt.data()); + de.Upload(eids); + Tensor tact = T2(da.ptr(), d, P, K); + Tensor twt = Tensor::Contiguous(dwt.ptr(), f.dt, d, {E * N, K}); + Tensor te = TI32(de.ptr(), d, P); + Tensor tout = T2(dout.ptr(), d, P, N); + vt::MatmulBTQuantGrouped(q, tout, tact, twt, te); + CHECK(Nmse(ref, dout.Download()) <= kNmseTol); + 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;