fix(BACKEND-ROCM): narrow silu(gate) to the gate dtype before multiply

.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<packed_t>`.
+- 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 `<typename Tin, typename Tout>`. The gate dtype is
+`Tin`. After computing `g / (1.0f + expf(-g))`, narrow via
+`NarrowTo<Tin>(...)` before multiplying by `up`:
+
+```cpp
+const float silu = NarrowTo<Tin>(g / (1.0f + expf(-g)));
+St(out, idx, silu * up);
+```
+
+### MoE `MoeSiluMulK` fix
+
+The kernel is templated `<typename Tout, typename Tg, typename Tu>`. The gate
+dtype is `Tg`. After computing `Silu(Ld(gate, i))`, narrow via
+`NarrowTo<Tg>(...)` before multiplying by `Ld(up, i)`:
+
+```cpp
+St(out, i, NarrowTo<Tg>(Silu(Ld(gate, i))) * Ld(up, i));
+```
+
+## Risks
+
+- **f32 paths unchanged**: `NarrowTo<float>` 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<Tin, Tout>` and `MoeSiluMulK<Tout, Tg, Tu>` —
+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.

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 <typename T>
+__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 <typename Tin, typename Tout>
 __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 <typename Tin, typename Tout>
 __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<Tin>(0.5f * g * (1.0f + tanhf(inner))) * up);
   }
 }
 
@@ -91,18 +109,22 @@ __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<Tin>(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 <typename Tin, typename Tout>
 __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;
        idx += gridDim.x * blockDim.x) {
     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<Tin>(g / (1.0f + expf(-g))) * up);
   }
 }
 

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 <typename T>
+__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 <typename Tout, typename Tg, typename Tu>
 __global__ void MoeSiluMulK(Tout* out, const Tg* gate, const Tu* up, int64_t n) {
   const int64_t step = static_cast<int64_t>(gridDim.x) * blockDim.x;
   for (int64_t i = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
        i < n; i += step) {
-    St(out, i, Silu(Ld(gate, i)) * Ld(up, i));
+    St(out, i, NarrowTo<Tg>(Silu(Ld(gate, i))) * Ld(up, i));
   }
 }
 

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