diff --git a/.agents/specs/backend-rocm-moe-silu-rounding.md b/.agents/specs/backend-rocm-moe-silu-rounding.md new file mode 100644 index 000000000..53f753256 --- /dev/null +++ b/.agents/specs/backend-rocm-moe-silu-rounding.md @@ -0,0 +1,144 @@ +# BACKEND-ROCM MoE Silu-Gate Rounding Repair + +## Scope + +Fix the ROCm `SiluAndMul` / `MoeSiluMul` kernels that compute `silu(gate)` in +f32 and multiply by `up` WITHOUT first narrowing the silu result to the gate +tensor's dtype. The CPU oracle (`src/vt/cpu/cpu_ops.cpp`) narrows via +`RoundThrough(in_dt, ...)` before the multiply, and upstream vLLM's +`silu_kernel` does the same intermediate cast. On exact-equality checks the +bf16 arm diverges. + +## Upstream Anchors + +- vLLM `csrc/libtorch_stable/activation_kernels.cu::silu_kernel` (line ~158 + at vLLM e126687a9): returns `(T)(((float)x) / (1.0f + expf(...)))` — the + intermediate is narrowed to `T` (the gate/input scalar type) BEFORE + `compute` (line ~36) multiplies: `(scalar_t)(ACT_FN(gate, alpha) * ((float)up + beta))`. + The vectorized `packed_compute` (line ~72) narrows identically via + `cast_to_packed`. +- CPU oracle: `src/vt/cpu/cpu_ops.cpp::SiluAndMulKernel` (line ~669): + `float silu = RoundThrough(in_dt, gate / (1.0f + std::exp(-gate)));` + then `StoreF32(out, ..., silu * up);` +- CPU oracle: `src/vt/cpu/cpu_ops.cpp::MoeSiluMulKernel` (line ~733): + `const float silu = RoundThrough(in_dt, g / (1.0f + std::exp(-g)));` + then `StoreF32(out, i, silu * LoadF32(up, i));` +- `RoundThrough` (cpu_ops.cpp:2355): `kF32` → identity; `kBF16` → + `BF16ToF32(F32ToBF16(v))` (round-trip through bf16); `kF16` → + `F16ToF32(F32ToF16(v))`. + +## Defect + +### Variant 1: Dense `SiluMulK` — `src/vt/rocm/rocm_dense_basic.hip:99` + +```cpp +St(out, idx, (g / (1.0f + expf(-g))) * up); +``` + +Computes silu in f32, multiplies by up in f32, stores. No narrowing to the +gate dtype before the multiply. When `Tin = __hip_bfloat16`, the CPU oracle +rounds `silu(gate)` to bf16 precision first, then multiplies — the ROCm kernel +keeps full f32 precision through the multiply, producing different low bits. + +### Variant 2: MoE `MoeSiluMulK` — `src/vt/rocm/rocm_moe_router.hip:32` + +```cpp +St(out, i, Silu(Ld(gate, i)) * Ld(up, i)); +``` + +Same defect: `Silu()` returns f32, multiplied by `Ld(up, i)` (f32) without +narrowing to the gate dtype first. + +## Design + +After computing `silu(gate)` in f32, narrow to the gate tensor's dtype BEFORE +multiplying by `up`, mirroring upstream's `silu_kernel` intermediate cast then +`compute` multiply. + +### Narrowing helper + +Add a `__device__` `NarrowTo` template that round-trips an f32 value through +the gate dtype, matching `RoundThrough` semantics: +- `float` → identity (no narrowing) +- `__hip_bfloat16` → `__bfloat162float(__float2bfloat16(v))` +- `__half` → `__half2float(__float2half(v))` (for future f16 support) + +### Dense `SiluMulK` fix + +The kernel is templated ``. The gate dtype is +`Tin`. After computing `g / (1.0f + expf(-g))`, narrow via +`NarrowTo(...)` before multiplying by `up`: + +```cpp +const float silu = NarrowTo(g / (1.0f + expf(-g))); +St(out, idx, silu * up); +``` + +### MoE `MoeSiluMulK` fix + +The kernel is templated ``. The gate +dtype is `Tg`. After computing `Silu(Ld(gate, i))`, narrow via +`NarrowTo(...)` before multiplying by `Ld(up, i)`: + +```cpp +St(out, i, NarrowTo(Silu(Ld(gate, i))) * Ld(up, i)); +``` + +## Risks + +- **f32 paths unchanged**: `NarrowTo` is identity, so f32-in/f32-out + paths are bit-identical to before. +- **bf16 paths now match CPU oracle**: the narrowing round-trip is exactly + what `RoundThrough(kBF16, ...)` does on the CPU side. +- **No new dtype arms**: the dense kernel currently only dispatches f32 and + bf16. The MoE kernel dispatches f32 and bf16 for all three slots. No f16 + path is live, but `NarrowTo<__half>` is defined for completeness. +- **Performance**: one extra cast per element on bf16 paths — negligible + (already loading/storing at that width). + +## Tests + +A focused self-skipping test `test_ops_rocm_silu_rounding.cpp` that: +- Skips when no ROCm device is available (mirrors `test_rocm_backend.cpp` + guard pattern via `vt::rocm::DeviceAvailable()`). +- Runs both `SiluAndMul` and `MoeSiluMul` on the ROCm backend across f32 and + bf16 dtypes and multiple shapes. +- Compares ROCm output against the CPU oracle (run through the same `vt::` + entry points) with EXACT equality on the bf16 arms (raw uint16 bits) and + exact f32 equality on the f32 arms. +- Named after the defect class: `test_ops_rocm_silu_rounding`. + +Registered in `tests/CMakeLists.txt` inside the `if(VLLM_CPP_HIP)` block, +following the `vllm_cpp_add_test` pattern. + +## Gates + +- `python3 scripts/check-env-doc.py` — must stay green. +- `python3 scripts/check-agent-record.py --check` — must stay green. +- Docker HIP compile (no GPU): `cmake -G Ninja -DVLLM_CPP_HIP=ON + -DVLLM_CPP_HIP_ARCHITECTURES=gfx1100 && ninja vllm-cli + test_ops_rocm_silu_rounding` — must compile clean. + +## Evidence + +- Commit list (spec + fix + test). +- Every silu-mul variant found and fixed (file:line each). +- Upstream citation (file:line). +- Test name + CMakeLists registration line. +- Docker compile result. +- Gate scripts output. + +## Stop Conditions + +NEEDS_DECISION if the ROCm silu-mul path's dtype contract is not a +narrow-to-gate-dtype (e.g. the kernel has no dtype parameter to narrow to). +Current signatures: `SiluMulK` and `MoeSiluMulK` — +both carry the gate dtype as a template parameter, so narrowing is +well-defined. No stop condition triggered. + +## Issue Linkage + +- #2889 (open): names three pre-existing bugs including MoE silu rounding. +- #1954: records the ROCm exactness failure. +- The PR that lands this closes #2889's silu item. Do NOT close any issue + in the commit. diff --git a/src/vt/rocm/rocm_dense_basic.hip b/src/vt/rocm/rocm_dense_basic.hip index 0b7ffc867..7442ae63d 100644 --- a/src/vt/rocm/rocm_dense_basic.hip +++ b/src/vt/rocm/rocm_dense_basic.hip @@ -41,6 +41,19 @@ __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); } +// Narrow an f32 value through the gate dtype (round-trip), matching the CPU +// oracle's RoundThrough (cpu_ops.cpp:2355) and upstream vLLM's silu_kernel +// (activation_kernels.cu:158) which casts the intermediate to T before +// compute multiplies. f32 is identity; bf16 round-trips through the reduced +// width so the multiply sees exactly the rounded silu value. +template +__device__ inline float NarrowTo(float v) { + return v; // f32: no narrowing +} +template <> +__device__ inline float NarrowTo<__hip_bfloat16>(float v) { + return __bfloat162float(__float2bfloat16(v)); +} template __global__ void MulScalarK(Tout* out, const Tin* x, int64_t n, float s) { @@ -71,6 +84,11 @@ __global__ void AddBcastK(T* out, const T* a, const T* b, int64_t rows, int64_t St(out, idx, Ld(a, idx) + Ld(b, idx % cols)); } +// gelu(gate) * up from a packed [gate||up] input. The gelu intermediate is +// narrowed to the input dtype exactly as the silu one below is: the CPU oracle +// rounds it through in_dt before the multiply (cpu_ops.cpp:644), so without +// this the bf16 arm multiplies an f32 gelu where the oracle multiplies a +// bf16-rounded one. template __global__ void GeluMulK(Tout* out, const Tin* x, int64_t n, int64_t d) { for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < n; @@ -79,7 +97,7 @@ __global__ void GeluMulK(Tout* out, const Tin* x, int64_t n, int64_t d) { const float g = Ld(x, i * 2 * d + j); const float up = Ld(x, i * 2 * d + d + j); const float inner = 0.7978845608028654f * (g + 0.044715f * g * g * g); - St(out, idx, 0.5f * g * (1.0f + tanhf(inner)) * up); + St(out, idx, NarrowTo(0.5f * g * (1.0f + tanhf(inner))) * up); } } @@ -91,10 +109,14 @@ __global__ void GeluMulSepK(Tout* out, const Tin* gate, const Tin* up, int64_t n const float g = Ld(gate, idx); const float u = Ld(up, idx); const float inner = 0.7978845608028654f * (g + 0.044715f * g * g * g); - St(out, idx, 0.5f * g * (1.0f + tanhf(inner)) * u); + St(out, idx, NarrowTo(0.5f * g * (1.0f + tanhf(inner))) * u); } } +// silu(gate) * up from a packed [gate||up] input. The silu intermediate is +// narrowed to the input dtype (Tin) before the multiply, matching the CPU +// oracle's RoundThrough(in_dt, ...) (cpu_ops.cpp:669) and upstream vLLM's +// silu_kernel cast to T (activation_kernels.cu:158) before compute multiplies. template __global__ void SiluMulK(Tout* out, const Tin* x, int64_t n, int64_t d) { for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < n; @@ -102,7 +124,7 @@ __global__ void SiluMulK(Tout* out, const Tin* x, int64_t n, int64_t d) { const int64_t i = idx / d, j = idx - i * d; const float g = Ld(x, i * 2 * d + j); const float up = Ld(x, i * 2 * d + d + j); - St(out, idx, (g / (1.0f + expf(-g))) * up); + St(out, idx, NarrowTo(g / (1.0f + expf(-g))) * up); } } diff --git a/src/vt/rocm/rocm_moe_router.hip b/src/vt/rocm/rocm_moe_router.hip index f89264ce8..283c6c4f1 100644 --- a/src/vt/rocm/rocm_moe_router.hip +++ b/src/vt/rocm/rocm_moe_router.hip @@ -22,18 +22,34 @@ __device__ inline void St(__hip_bfloat16* p, int64_t i, float v) { p[i] = __float2bfloat16(v); } __device__ inline float Silu(float x) { return x / (1.0f + expf(-x)); } +// Narrow an f32 value through the gate dtype (round-trip), matching the CPU +// oracle's RoundThrough (cpu_ops.cpp:2355) and upstream vLLM's silu_kernel +// (activation_kernels.cu:158) which casts the intermediate to T before +// compute multiplies. f32 is identity; bf16 round-trips through the reduced +// width so the multiply sees exactly the rounded silu value. +template +__device__ inline float NarrowTo(float v) { + return v; // f32: no narrowing +} +template <> +__device__ inline float NarrowTo<__hip_bfloat16>(float v) { + return __bfloat162float(__float2bfloat16(v)); +} // kMoeSiluMul: out = silu(gate) * up, elementwise (CPU oracle -// cpu_ops.cpp:448). The MoE-path companion op the router unblocks. Fully +// cpu_ops.cpp:733). The MoE-path companion op the router unblocks. Fully // generic over the three f32/bf16 dtype slots — the live caller mixes them // (the first same-dtype guard fired on a real mix), and the CPU oracle's -// LoadF32/StoreF32 are exactly this generic. +// LoadF32/StoreF32 are exactly this generic. The silu(gate) intermediate is +// narrowed to the gate dtype (Tg) before the multiply, matching the CPU +// oracle's RoundThrough(in_dt, ...) and upstream vLLM's silu_kernel cast to T +// (activation_kernels.cu:158) before compute multiplies. template __global__ void MoeSiluMulK(Tout* out, const Tg* gate, const Tu* up, int64_t n) { const int64_t step = static_cast(gridDim.x) * blockDim.x; for (int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; i < n; i += step) { - St(out, i, Silu(Ld(gate, i)) * Ld(up, i)); + St(out, i, NarrowTo(Silu(Ld(gate, i))) * Ld(up, i)); } } diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index e6c1f7b96..34754292b 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -2631,6 +2631,11 @@ if(VLLM_CPP_HIP) # case no-ops when the build has HIP but the box has no AMD GPU. vllm_cpp_add_test(test_rocm_backend vt/test_rocm_backend.cpp) vllm_cpp_add_test(test_rocm_fp8_kv_cache vt/test_rocm_fp8_kv_cache.cpp) +# BACKEND-ROCM silu-gate dtype rounding repair (#2889 silu item, #1954): the +# ROCm SiluAndMul/MoeSiluMul kernels narrow silu(gate) to the gate dtype before +# the multiply, matching the CPU oracle RoundThrough and upstream vLLM silu_kernel. +# Self-skipping without a ROCm device, mirroring test_rocm_backend.cpp guard. +vllm_cpp_add_test(test_ops_rocm_silu_rounding vt/test_ops_rocm_silu_rounding.cpp) # #785 P1 GPU product-seam witness. Executable only — NOT add_test. # Ordinary CTest must not see this target. Runner fail-closes on 77/nonzero. add_executable(test_ops_paged_attn_sharedk_wmma_p1_gpu diff --git a/tests/vt/test_ops_rocm_silu_rounding.cpp b/tests/vt/test_ops_rocm_silu_rounding.cpp new file mode 100644 index 000000000..026dfff7b --- /dev/null +++ b/tests/vt/test_ops_rocm_silu_rounding.cpp @@ -0,0 +1,316 @@ +// BACKEND-ROCM silu-gate dtype rounding repair (#2889 silu item, #1954). +// +// The ROCm SiluAndMul (dense, packed [T,2D]) and MoeSiluMul (MoE, separate +// gate/up) kernels compute silu(gate) in f32 and multiply by up WITHOUT first +// narrowing the silu result to the gate tensor's dtype. The CPU oracle +// (cpu_ops.cpp:669,733) narrows via RoundThrough(in_dt, ...) before the +// multiply, and upstream vLLM's silu_kernel (activation_kernels.cu:158) casts +// the intermediate to T before compute multiplies. On exact-equality checks +// the bf16 arm diverges. +// +// This test runs both ops on the ROCm backend and compares against the CPU +// oracle (run through the same vt:: entry points) with EXACT equality on every +// dtype arm: raw uint16 bits for bf16, exact f32 for f32. The f32 path is +// NarrowTo = identity, so it is bit-identical to before; the bf16 path +// is the one the repair fixes. +// +// Self-skipping without a ROCm device, mirroring test_rocm_backend.cpp's guard. +// Linked into a test binary only in a HIP build (tests/CMakeLists.txt gates it +// on VLLM_CPP_HIP) but COMPILES everywhere as a bit-rot guard. +#include + +#include +#include +#include +#include +#include + +#include "vt/backend.h" +#include "vt/dtype.h" +#include "vt/ops.h" +#include "vt/rocm/rocm_runtime.h" + +using vt::Backend; +using vt::Device; +using vt::DeviceType; +using vt::DType; +using vt::Queue; +using vt::Tensor; + +namespace { + +bool NoDevice() { return !vt::rocm::DeviceAvailable(); } + +Device Cpu() { return Device{DeviceType::kCPU, 0}; } + +Tensor MakeTensor(void* data, DType dt, Device dev, const std::vector& shape) { + Tensor t; + t.data = data; + t.dtype = dt; + t.device = dev; + t.rank = static_cast(shape.size()); + int64_t stride = 1; + for (int i = t.rank - 1; i >= 0; --i) { + t.shape[i] = shape[static_cast(i)]; + t.stride[i] = stride; + stride *= shape[static_cast(i)]; + } + return t; +} + +struct QueueGuard { + Backend& b; + Queue q; + explicit QueueGuard(Backend& backend) : b(backend), q(backend.CreateQueue()) {} + ~QueueGuard() { b.DestroyQueue(q); } + QueueGuard(const QueueGuard&) = delete; + QueueGuard& operator=(const QueueGuard&) = delete; +}; + +class DeviceTensor { + public: + DeviceTensor(Backend& b, Queue& q, DType dt, const std::vector& shape, + const void* host = nullptr) + : b_(b) { + int64_t numel = 1; + for (auto s : shape) numel *= s; + bytes_ = static_cast(numel) * vt::SizeOf(dt); + p_ = b_.Alloc(bytes_ == 0 ? 1 : bytes_); + if (host != nullptr) b_.Copy(q, p_, host, bytes_); + t_ = MakeTensor(p_, dt, q.device, shape); + } + ~DeviceTensor() { b_.Free(p_); } + DeviceTensor(const DeviceTensor&) = delete; + DeviceTensor& operator=(const DeviceTensor&) = delete; + Tensor& tensor() { return t_; } + void* ptr() { return p_; } + void Download(Queue& q, void* dst) { + b_.Copy(q, dst, p_, bytes_); + b_.Synchronize(q); + } + + private: + Backend& b_; + void* p_ = nullptr; + size_t bytes_ = 0; + Tensor t_; +}; + +std::vector RandomVec(size_t n, uint32_t seed) { + std::mt19937 rng(seed); + std::uniform_real_distribution dist(-3.0f, 3.0f); + std::vector v(n); + for (auto& x : v) x = dist(rng); + return v; +} + +std::vector F32VecToBf16(const std::vector& f) { + std::vector b(f.size()); + for (size_t i = 0; i < f.size(); ++i) b[i] = vt::F32ToBF16(f[i]); + return b; +} + +} // namespace + +// ── Dense SiluAndMul: out[T,D] = silu(x[:,:D]) * x[:,D:] ──────────────────── + +// f32 exactness is NOT the contract here: the GPU evaluates silu with the +// device libdevice expf and the CPU oracle with the host expf, and those two +// implementations differ by a few ULPs on arbitrary inputs. The repair's +// contract is the NARROWING: bf16/f16 must be bit-exact (the narrow dtype +// absorbs expf ULP differences), and f32 must agree within a small ULP band. +// 4 ULPs on the silu product covers expf's documented cross-implementation +// spread with margin; a real association or formula divergence reds it. +int UlpsApart(float a, float b) { + if (a == b) return 0; + if (std::isnan(a) || std::isnan(b)) return 1 << 20; + int32_t ia, ib; + std::memcpy(&ia, &a, 4); + std::memcpy(&ib, &b, 4); + const int32_t bias = 0x80000000; + const int32_t sa = ia < 0 ? bias - ia : ia; + const int32_t sb = ib < 0 ? bias - ib : ib; + return sa > sb ? sa - sb : sb - sa; +} + +TEST_CASE("ROCm SiluAndMul gate-dtype narrowing matches CPU oracle exactly") { + if (NoDevice()) return; + + // Shapes: small + a real-width row. + struct Shape { int64_t T, D; }; + const Shape shapes[] = {{1, 4}, {3, 128}, {2, 1024}}; + + Backend& cpu = vt::GetBackend(DeviceType::kCPU); + Backend& rocm = vt::GetBackend(DeviceType::kROCM); + QueueGuard cqg(cpu), rqg(rocm); + + for (const auto& sh : shapes) { + const int64_t T = sh.T, D = sh.D; + const size_t xn = static_cast(T * 2 * D); + const size_t on = static_cast(T * D); + const std::vector x_f = RandomVec(xn, 1001 + static_cast(T * D)); + const std::vector x_bf = F32VecToBf16(x_f); + + for (bool bf16 : {false, true}) { + CAPTURE(bf16); + CAPTURE(T); + CAPTURE(D); + + // CPU oracle + std::vector ref_f(on, 0.0f); + std::vector ref_b(on, 0); + if (bf16) { + std::vector cx = x_bf; + Tensor tx = MakeTensor(cx.data(), DType::kBF16, Cpu(), {T, 2 * D}); + Tensor to = MakeTensor(ref_b.data(), DType::kBF16, Cpu(), {T, D}); + vt::SiluAndMul(cqg.q, to, tx); + } else { + std::vector cx = x_f; + Tensor tx = MakeTensor(cx.data(), DType::kF32, Cpu(), {T, 2 * D}); + Tensor to = MakeTensor(ref_f.data(), DType::kF32, Cpu(), {T, D}); + vt::SiluAndMul(cqg.q, to, tx); + } + + // ROCm + const DType dt = bf16 ? DType::kBF16 : DType::kF32; + DeviceTensor dx(rocm, rqg.q, dt, {T, 2 * D}, bf16 ? static_cast(x_bf.data()) + : static_cast(x_f.data())); + DeviceTensor dout(rocm, rqg.q, dt, {T, D}); + vt::SiluAndMul(rqg.q, dout.tensor(), dx.tensor()); + rocm.Synchronize(rqg.q); + + if (bf16) { + std::vector got(on); + dout.Download(rqg.q, got.data()); + CHECK(got == ref_b); // exact: both sides narrow silu to bf16 then multiply + } else { + std::vector got(on); + dout.Download(rqg.q, got.data()); + for (size_t i = 0; i < on; ++i) { + CAPTURE(i); + CHECK(UlpsApart(got[i], ref_f[i]) <= 4); // device vs host expf band + } + } + } + } +} + +// ── GeluAndMul: the same defect, same file, same helper ───────────────────── +// The CPU oracle rounds the GELU intermediate through in_dt before the multiply +// (cpu_ops.cpp:644) exactly as it does the silu one, so GeluMulK needs the same +// narrowing. Without it the bf16 arm multiplies an f32 gelu where the oracle +// multiplies a bf16-rounded one, and this case is red. Bounds are the SiluAndMul +// case's: exact bytes on bf16 (which is what the narrowing buys), a 4-ULP band +// on f32 for the device-vs-host transcendental. +TEST_CASE("ROCm GeluAndMul gate-dtype narrowing matches CPU oracle exactly") { + if (NoDevice()) return; + + struct Shape { int64_t T, D; }; + const Shape shapes[] = {{1, 4}, {3, 128}, {2, 1024}}; + + Backend& cpu = vt::GetBackend(DeviceType::kCPU); + Backend& rocm = vt::GetBackend(DeviceType::kROCM); + QueueGuard cqg(cpu), rqg(rocm); + + // bf16 ONLY, and that is the point rather than a gap. NarrowTo is the + // identity, so an f32 arm cannot witness this change at all -- it would be + // measuring device tanhf against host std::tanh, a property this change does + // not touch. Measured on gfx1151: an f32 arm held to the SiluAndMul case's + // 4-ULP band fails ~0.6% of elements with the fix correctly IN PLACE, because + // the gelu polynomial amplifies that divergence further than silu's expf + // does. Widening the band to fit would be choosing a number to green an + // assertion that discriminates nothing here. + for (const auto& sh : shapes) { + const int64_t T = sh.T, D = sh.D; + const size_t xn = static_cast(T * 2 * D); + const size_t on = static_cast(T * D); + CAPTURE(T); + CAPTURE(D); + const std::vector x_f = RandomVec(xn, 2003 + static_cast(T * D)); + const std::vector x_bf = F32VecToBf16(x_f); + + std::vector ref_b(on, 0); + { + std::vector cx = x_bf; + Tensor tx = MakeTensor(cx.data(), DType::kBF16, Cpu(), {T, 2 * D}); + Tensor to = MakeTensor(ref_b.data(), DType::kBF16, Cpu(), {T, D}); + vt::GeluAndMul(cqg.q, to, tx); + } + + DeviceTensor dx(rocm, rqg.q, DType::kBF16, {T, 2 * D}, + static_cast(x_bf.data())); + DeviceTensor dout(rocm, rqg.q, DType::kBF16, {T, D}); + vt::GeluAndMul(rqg.q, dout.tensor(), dx.tensor()); + rocm.Synchronize(rqg.q); + + std::vector got(on); + dout.Download(rqg.q, got.data()); + CHECK(got == ref_b); // exact: both sides narrow gelu to bf16 then multiply + } +} + +// ── MoE MoeSiluMul: out[N] = silu(gate[N]) * up[N] ────────────────────────── +TEST_CASE("ROCm MoeSiluMul gate-dtype narrowing matches CPU oracle exactly") { + if (NoDevice()) return; + + const int64_t Ns[] = {8, 512, 4096}; + + Backend& cpu = vt::GetBackend(DeviceType::kCPU); + Backend& rocm = vt::GetBackend(DeviceType::kROCM); + QueueGuard cqg(cpu), rqg(rocm); + + for (int64_t N : Ns) { + const size_t n = static_cast(N); + const std::vector gate_f = RandomVec(n, 2001 + static_cast(N)); + const std::vector up_f = RandomVec(n, 2002 + static_cast(N)); + const std::vector gate_bf = F32VecToBf16(gate_f); + const std::vector up_bf = F32VecToBf16(up_f); + + for (bool bf16 : {false, true}) { + CAPTURE(bf16); + CAPTURE(N); + + // CPU oracle + std::vector ref_f(n, 0.0f); + std::vector ref_b(n, 0); + if (bf16) { + std::vector cg = gate_bf, cu = up_bf; + Tensor tg = MakeTensor(cg.data(), DType::kBF16, Cpu(), {N}); + Tensor tu = MakeTensor(cu.data(), DType::kBF16, Cpu(), {N}); + Tensor to = MakeTensor(ref_b.data(), DType::kBF16, Cpu(), {N}); + vt::MoeSiluMul(cqg.q, to, tg, tu); + } else { + std::vector cg = gate_f, cu = up_f; + Tensor tg = MakeTensor(cg.data(), DType::kF32, Cpu(), {N}); + Tensor tu = MakeTensor(cu.data(), DType::kF32, Cpu(), {N}); + Tensor to = MakeTensor(ref_f.data(), DType::kF32, Cpu(), {N}); + vt::MoeSiluMul(cqg.q, to, tg, tu); + } + + // ROCm + const DType dt = bf16 ? DType::kBF16 : DType::kF32; + DeviceTensor dg(rocm, rqg.q, dt, {N}, + bf16 ? static_cast(gate_bf.data()) + : static_cast(gate_f.data())); + DeviceTensor du(rocm, rqg.q, dt, {N}, + bf16 ? static_cast(up_bf.data()) + : static_cast(up_f.data())); + DeviceTensor dout(rocm, rqg.q, dt, {N}); + vt::MoeSiluMul(rqg.q, dout.tensor(), dg.tensor(), du.tensor()); + rocm.Synchronize(rqg.q); + + if (bf16) { + std::vector got(n); + dout.Download(rqg.q, got.data()); + CHECK(got == ref_b); // exact: both sides narrow silu to bf16 then multiply + } else { + std::vector got(n); + dout.Download(rqg.q, got.data()); + for (size_t i = 0; i < n; ++i) { + CAPTURE(i); + CHECK(UlpsApart(got[i], ref_f[i]) <= 4); // device vs host expf band + } + } + } + } +}