rocm: fused flash-attention DECODE (single-query, strided-KV, GQA)

hanzo-ml 0.11.83 + hanzo-nn 0.11.39 + hanzo-rocm-kernels 0.11.35.

ROCm single-query decode had no fused attention: it fell to eager repeat_kv
(Tensor::cat -> copy2d) + QKᵀ bmm + f32-softmax + P@V bmm + a per-layer
.contiguous() copy of the WHOLE active KV cache -- ~10 dispatches and the 653
copyBufferRectAligned/token that dominated ROCm decode glue.

flash_decode_{f16,bf16} (flash.hip): one block per (batch, q_head), 128 threads,
GQA implicit (kv_head = q_head/(Hq/Hkv), no repeat_kv), KV cache read IN PLACE via
(batch,head,seq) strides (head dim contiguous, no cache copy), tiled online softmax
in f32. RocmStorage::flash_attn_decode launches it from strided k/v layouts;
hanzo_nn::attention::rocm_flash_attn_decode is the Tensor wrapper; the engine wires
it as rocm_decode_attn (mirrors vulkan_decode_attn's gating).

Correct vs an f64 eager oracle across GQA ratios, tile-boundary L, and the strided
cache layout (test flash_decode_matches_eager, rel 2.4e-4..3.9e-4). copyBufferRect
33950->5760/run (-83%); decode QKᵀ/softmax bmms 3708->216 (prefill only).

Cargo.toml

@@ -47,14 +47,14 @@ license = "MIT OR Apache-2.0"
 
 [workspace.dependencies]
 rocm-rs = { package = "hanzo-rocm", version = "0.5.2", default-features = false, features = ["macros"] }
-hanzo-rocm-kernels = { path = "./hanzo-rocm-kernels", version = "0.11.34" }
+hanzo-rocm-kernels = { path = "./hanzo-rocm-kernels", version = "0.11.35" }
 ab_glyph = "0.2.23"
 accelerate-src = { version = "0.3.2" }
 anyhow = { version = "1", features = ["backtrace"] }
 byteorder = "1.4.3"
 chrono = "0.4"
 # Hanzo-renamed workspace dependencies (at version 0.10.x)
-hanzo-ml = { path = "./hanzo-ml", package = "hanzo-ml", version = "0.11.38" }
+hanzo-ml = { path = "./hanzo-ml", package = "hanzo-ml", version = "0.11.83" }
 hanzo-datasets = { path = "./hanzo-datasets", version = "0.11.8" }
 hanzo-flash-attn = { path = "./hanzo-flash-attn", version = "0.11.34" }
 hanzo-flash-attn-v3 = { path = "./hanzo-flash-attn-v3", version = "0.11.8" }

hanzo-ml/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "hanzo-ml"
-version = "0.11.82"
+version = "0.11.83"
 edition = "2021"
 description = "Fast multi-backend tensor & ML framework for Rust (CPU/CUDA/Metal/Vulkan/ROCm) with quantization — the compute core of the Hanzo stack."
 repository = "https://github.com/hanzoai/ml"

hanzo-ml/src/rocm_backend/mod.rs

@@ -3024,6 +3024,127 @@ impl RocmStorage {
         }
     }
 
+    /// Fused flash-attention DECODE: single query per head (`Lq == 1`) attending the whole KV cache.
+    /// `q` is contiguous `[B, Hq, 1, D]`; `k`/`v` are the KV cache read IN PLACE via their (batch,
+    /// head, seq) strides (innermost head-dim contiguous), so there is no repeat_kv copy and no
+    /// whole-cache `.contiguous()`. GQA via `Hq % Hkv`, `D == 128`, f16/bf16. One block per (batch,
+    /// q_head); f32 online softmax. Output `[B, Hq, 1, D]` matches eager naive_sdpa within f16 rounding.
+    pub fn flash_attn_decode(
+        &self,
+        q_layout: &Layout,
+        k: &RocmStorage,
+        k_layout: &Layout,
+        v: &RocmStorage,
+        v_layout: &Layout,
+        scale: f32,
+    ) -> Result<Self> {
+        use hanzo_rocm_kernels::kernel::{FlashKernel, KernelSource};
+        const FD_DH: usize = 128;
+        let q = self;
+        if !q_layout.is_contiguous() {
+            crate::bail!("flash_attn_decode requires contiguous q");
+        }
+        let (b, hq, lq, d) = q_layout.shape().dims4()?;
+        let (kb, hkv, lk, kd) = k_layout.shape().dims4()?;
+        let (vb, vhkv, vlk, vd) = v_layout.shape().dims4()?;
+        if lq != 1 {
+            crate::bail!("flash_attn_decode requires q seq len 1, got {lq}");
+        }
+        if d != FD_DH || kd != FD_DH || vd != FD_DH {
+            crate::bail!("flash_attn_decode requires head_dim {FD_DH}, got q={d} k={kd} v={vd}");
+        }
+        if (kb, vb, vhkv, vlk) != (b, b, hkv, lk) {
+            crate::bail!(
+                "flash_attn_decode q/k/v batch/kv-head/kv-len mismatch {k_layout:?} {v_layout:?}"
+            );
+        }
+        if hkv == 0 || hq % hkv != 0 {
+            crate::bail!("flash_attn_decode requires Hq % Hkv == 0, got Hq={hq} Hkv={hkv}");
+        }
+        let ks = k_layout.stride();
+        let vs = v_layout.stride();
+        if ks[3] != 1 || vs[3] != 1 {
+            crate::bail!("flash_attn_decode requires innermost (head-dim) stride 1");
+        }
+        let device = self.device.clone();
+        let o_el = b * hq * d; // lq == 1
+        let (b_i, hq_i, hkv_i, lk_i) = (b as i32, hq as i32, hkv as i32, lk as i32);
+        let (kb_s, kh_s, kl_s) = (ks[0] as i64, ks[1] as i64, ks[2] as i64);
+        let (vb_s, vh_s, vl_s) = (vs[0] as i64, vs[1] as i64, vs[2] as i64);
+        const BLOCK: u32 = 128;
+        let grid = rocm_rs::hip::Dim3 {
+            x: hq as u32,
+            y: b as u32,
+            z: 1,
+        };
+        let block = rocm_rs::hip::Dim3::from(BLOCK);
+        let q_off = q_layout.start_offset();
+        let k_off = k_layout.start_offset();
+        let v_off = v_layout.start_offset();
+
+        macro_rules! launch_decode {
+            ($variant:ident, $ty:ty, $func:literal) => {{
+                let q_mem = match &q.slice {
+                    RocmStorageSlice::$variant(m) => m,
+                    _ => unreachable!(),
+                };
+                let k_mem = match &k.slice {
+                    RocmStorageSlice::$variant(m) => m,
+                    _ => crate::bail!("flash_attn_decode: k dtype must match q dtype"),
+                };
+                let v_mem = match &v.slice {
+                    RocmStorageSlice::$variant(m) => m,
+                    _ => crate::bail!("flash_attn_decode: v dtype must match q dtype"),
+                };
+                let o_out = device.alloc::<$ty>(o_el)?;
+                let q_ptr = unsafe { q_mem.offset_ptr(q_off) };
+                let k_ptr = unsafe { k_mem.offset_ptr(k_off) };
+                let v_ptr = unsafe { v_mem.offset_ptr(v_off) };
+                let o_ptr = o_out.as_ptr();
+                unsafe {
+                    launch_kernel(
+                        &device,
+                        FlashKernel::NAME,
+                        FlashKernel::CODE,
+                        $func,
+                        grid,
+                        block,
+                        &mut [
+                            &b_i as *const i32 as *mut std::ffi::c_void,
+                            &hq_i as *const i32 as *mut std::ffi::c_void,
+                            &hkv_i as *const i32 as *mut std::ffi::c_void,
+                            &lk_i as *const i32 as *mut std::ffi::c_void,
+                            &scale as *const f32 as *mut std::ffi::c_void,
+                            (&q_ptr) as *const *mut std::ffi::c_void as *mut std::ffi::c_void,
+                            (&k_ptr) as *const *mut std::ffi::c_void as *mut std::ffi::c_void,
+                            &kb_s as *const i64 as *mut std::ffi::c_void,
+                            &kh_s as *const i64 as *mut std::ffi::c_void,
+                            &kl_s as *const i64 as *mut std::ffi::c_void,
+                            (&v_ptr) as *const *mut std::ffi::c_void as *mut std::ffi::c_void,
+                            &vb_s as *const i64 as *mut std::ffi::c_void,
+                            &vh_s as *const i64 as *mut std::ffi::c_void,
+                            &vl_s as *const i64 as *mut std::ffi::c_void,
+                            (&o_ptr) as *const *mut std::ffi::c_void as *mut std::ffi::c_void,
+                        ],
+                    )?;
+                }
+                Ok(Self {
+                    slice: RocmStorageSlice::$variant(o_out),
+                    device: device.clone(),
+                })
+            }};
+        }
+
+        match &q.slice {
+            RocmStorageSlice::F16(_) => launch_decode!(F16, f16, "flash_decode_f16"),
+            RocmStorageSlice::BF16(_) => launch_decode!(BF16, bf16, "flash_decode_bf16"),
+            other => crate::bail!(
+                "flash_attn_decode unsupported dtype {:?} (f16/bf16 only)",
+                other.dtype()
+            ),
+        }
+    }
+
     /// Fused softmax over the last dim (max-subtract-exp-sum-div, f32 accumulation), replacing the
     /// composite that ran 5 ops + 2 casts. One warp per row (32 lanes along threadIdx.y); the
     /// reduction is a warp shuffle (no shared memory). Reuses the ReduceKernel `softmax_{...}`.
@@ -5224,3 +5345,188 @@ mod dsl_norm_bench {
         }
     }
 }
+
+/// Fused flash-attention DECODE kernel vs an independent f64 eager reference (scale-relative gate).
+/// Covers GQA ratios, non-multiple-of-tile sequence lengths, and the STRIDED KV-cache layout the
+/// decode fast-path relies on (seq stride reflecting a padded max_len). This is the correctness gate
+/// for `rocm_decode_attn` -- the shared attention path across every model.
+#[cfg(all(test, feature = "rocm"))]
+mod flash_decode_test {
+    use super::*;
+    use half::f16;
+
+    fn rnd(n: usize, seed: u64) -> Vec<f16> {
+        let mut s = seed | 1;
+        (0..n)
+            .map(|_| {
+                s ^= s << 13;
+                s ^= s >> 7;
+                s ^= s << 17;
+                f16::from_f32(((s % 2000) as f32 / 1000.0 - 1.0) * 0.5)
+            })
+            .collect()
+    }
+
+    // f64 eager SDPA reference from the f16-rounded inputs (GQA): scores = scale*(Q.K), softmax, .V.
+    fn ref_out(
+        qf: &[f16],
+        kf: &[f16],
+        vf: &[f16],
+        b: usize,
+        hq: usize,
+        hkv: usize,
+        l: usize,
+        d: usize,
+        scale: f32,
+    ) -> Vec<f32> {
+        let gqa = hq / hkv;
+        let mut out = vec![0f32; b * hq * d];
+        for bi in 0..b {
+            for qh in 0..hq {
+                let kvh = qh / gqa;
+                let mut sc = vec![0f64; l];
+                let mut mx = f64::NEG_INFINITY;
+                for (j, scj) in sc.iter_mut().enumerate() {
+                    let mut s = 0f64;
+                    for t in 0..d {
+                        s += qf[(bi * hq + qh) * d + t].to_f64()
+                            * kf[((bi * hkv + kvh) * l + j) * d + t].to_f64();
+                    }
+                    *scj = s * scale as f64;
+                    if *scj > mx {
+                        mx = *scj;
+                    }
+                }
+                let mut den = 0f64;
+                for scj in sc.iter_mut() {
+                    *scj = (*scj - mx).exp();
+                    den += *scj;
+                }
+                for t in 0..d {
+                    let mut acc = 0f64;
+                    for (j, &p) in sc.iter().enumerate() {
+                        acc += p * vf[((bi * hkv + kvh) * l + j) * d + t].to_f64();
+                    }
+                    out[(bi * hq + qh) * d + t] = (acc / den) as f32;
+                }
+            }
+        }
+        out
+    }
+
+    fn max_rel(want: &[f32], got: &[f32]) -> f32 {
+        let mut maxabs = 0f32;
+        let mut maxref = 1e-6f32;
+        for i in 0..want.len() {
+            maxabs = maxabs.max((want[i] - got[i]).abs());
+            maxref = maxref.max(want[i].abs());
+        }
+        maxabs / maxref
+    }
+
+    // Pad the contiguous [b,hkv,l,d] KV to [b,hkv,maxl,d] so the layout's seq stride reflects maxl
+    // (the real decode-cache shape) -- exercises the kernel's strided read.
+    fn pad_seq(kf: &[f16], b: usize, hkv: usize, l: usize, d: usize, maxl: usize) -> Vec<f16> {
+        let mut out = vec![f16::ZERO; b * hkv * maxl * d];
+        for bi in 0..b {
+            for h in 0..hkv {
+                for j in 0..l {
+                    for t in 0..d {
+                        out[((bi * hkv + h) * maxl + j) * d + t] =
+                            kf[((bi * hkv + h) * l + j) * d + t];
+                    }
+                }
+            }
+        }
+        out
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    fn run(
+        device: &RocmDevice,
+        qf: &[f16],
+        kf: &[f16],
+        vf: &[f16],
+        b: usize,
+        hq: usize,
+        hkv: usize,
+        l: usize,
+        d: usize,
+        scale: f32,
+        maxl: Option<usize>,
+    ) -> Vec<f32> {
+        let q_st = RocmStorage {
+            slice: RocmStorageSlice::F16(device.clone_htod(qf).unwrap()),
+            device: device.clone(),
+        };
+        let k_st = RocmStorage {
+            slice: RocmStorageSlice::F16(device.clone_htod(kf).unwrap()),
+            device: device.clone(),
+        };
+        let v_st = RocmStorage {
+            slice: RocmStorageSlice::F16(device.clone_htod(vf).unwrap()),
+            device: device.clone(),
+        };
+        let q_l = Layout::contiguous((b, hq, 1, d));
+        let (k_l, v_l) = match maxl {
+            None => (
+                Layout::contiguous((b, hkv, l, d)),
+                Layout::contiguous((b, hkv, l, d)),
+            ),
+            Some(ml) => {
+                let st = vec![hkv * ml * d, ml * d, d, 1];
+                (
+                    Layout::new((b, hkv, l, d).into(), st.clone(), 0),
+                    Layout::new((b, hkv, l, d).into(), st, 0),
+                )
+            }
+        };
+        let out = q_st
+            .flash_attn_decode(&q_l, &k_st, &k_l, &v_st, &v_l, scale)
+            .unwrap();
+        let om = match out.slice {
+            RocmStorageSlice::F16(m) => m,
+            _ => panic!("expected f16 out"),
+        };
+        device
+            .clone_dtoh(&om)
+            .unwrap()
+            .iter()
+            .map(|x: &f16| x.to_f32())
+            .collect()
+    }
+
+    #[test]
+    fn flash_decode_matches_eager() {
+        let device = RocmDevice::new(0).unwrap();
+        let d = 128usize;
+        // (B, Hq, Hkv, L): MHA, GQA 4:1 and 8:1, tile-boundary and multi-tile L, batch>1.
+        for &(b, hq, hkv, l) in &[
+            (1, 8, 8, 64),
+            (1, 8, 2, 100),
+            (2, 16, 4, 37),
+            (1, 4, 4, 256),
+            (1, 14, 2, 200),
+            (1, 16, 16, 300),
+        ] {
+            let scale = 1.0f32 / (d as f32).sqrt();
+            let qf = rnd(b * hq * d, 0x11 + b as u64 * 7 + hq as u64);
+            let kf = rnd(b * hkv * l * d, 0x22 + l as u64 * 3 + hkv as u64);
+            let vf = rnd(b * hkv * l * d, 0x33 + l as u64 + hkv as u64 * 5);
+            let want = ref_out(&qf, &kf, &vf, b, hq, hkv, l, d, scale);
+
+            let got = run(&device, &qf, &kf, &vf, b, hq, hkv, l, d, scale, None);
+            let rel = max_rel(&want, &got);
+            println!("[flash_decode] B{b} Hq{hq} Hkv{hkv} L{l} contiguous rel={rel:.2e}");
+            assert!(rel < 1e-2, "flash_decode contiguous rel {rel} at {b}/{hq}/{hkv}/{l}");
+
+            let maxl = l + 17;
+            let kf_p = pad_seq(&kf, b, hkv, l, d, maxl);
+            let vf_p = pad_seq(&vf, b, hkv, l, d, maxl);
+            let got_s = run(&device, &qf, &kf_p, &vf_p, b, hq, hkv, l, d, scale, Some(maxl));
+            let rel_s = max_rel(&want, &got_s);
+            println!("[flash_decode] B{b} Hq{hq} Hkv{hkv} L{l} strided(maxl={maxl}) rel={rel_s:.2e}");
+            assert!(rel_s < 1e-2, "flash_decode strided rel {rel_s} at {b}/{hq}/{hkv}/{l}");
+        }
+    }
+}

hanzo-nn/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "hanzo-nn"
-version = "0.11.38"
+version = "0.11.39"
 edition = "2021"
 description = "Multi-backend tensor & ML framework for Rust — part of the Hanzo ML stack."
 repository = "https://github.com/hanzoai/ml"

hanzo-nn/src/attention/mod.rs

@@ -6,7 +6,7 @@ use hanzo_ml::Tensor;
 
 pub use cpu_flash::flash_attn;
 pub use cpu_flash::varlen::flash_attn_varlen_cpu;
-pub use rocm::rocm_flash_attn;
+pub use rocm::{rocm_flash_attn, rocm_flash_attn_decode};
 pub use varlen::flash_attn_varlen_unfused;
 
 #[derive(Debug, Clone, Default)]

hanzo-nn/src/attention/rocm.rs

@@ -59,3 +59,51 @@ pub fn rocm_flash_attn(
     let v = v.contiguous()?;
     q.apply_op3_no_bwd(&k, &v, &RocmFlashAttn { scale, causal })
 }
+
+/// Fused flash-attention DECODE (single query per head). Unlike `rocm_flash_attn`, `k`/`v` are read
+/// STRIDED (in place) so the caller skips the per-layer `.contiguous()` copy of the whole active KV
+/// cache. Inputs: `q` `[B, Hq, 1, 128]` (made contiguous here — a single token), `k`/`v` `[B, Hkv,
+/// Lk, 128]` with innermost head-dim contiguous, f16/bf16, same ROCm device. A lone decode query sees
+/// only past keys, so the whole cache is unmasked. Returns `[B, Hq, 1, 128]`.
+#[derive(Debug, Clone)]
+#[cfg_attr(not(feature = "rocm"), allow(dead_code))]
+struct RocmFlashDecode {
+    scale: f32,
+}
+
+impl hanzo_ml::CustomOp3 for RocmFlashDecode {
+    fn name(&self) -> &'static str {
+        "rocm-flash-decode"
+    }
+
+    fn cpu_fwd(
+        &self,
+        _s1: &hanzo_ml::CpuStorage,
+        _l1: &Layout,
+        _s2: &hanzo_ml::CpuStorage,
+        _l2: &Layout,
+        _s3: &hanzo_ml::CpuStorage,
+        _l3: &Layout,
+    ) -> Result<(hanzo_ml::CpuStorage, Shape)> {
+        hanzo_ml::bail!("rocm-flash-decode has no CPU path; route to naive_sdpa off ROCm")
+    }
+
+    #[cfg(feature = "rocm")]
+    fn rocm_fwd(
+        &self,
+        s1: &hanzo_ml::RocmStorage,
+        l1: &Layout,
+        s2: &hanzo_ml::RocmStorage,
+        l2: &Layout,
+        s3: &hanzo_ml::RocmStorage,
+        l3: &Layout,
+    ) -> Result<(hanzo_ml::RocmStorage, Shape)> {
+        let out = s1.flash_attn_decode(l1, s2, l2, s3, l3, self.scale)?;
+        Ok((out, l1.shape().clone()))
+    }
+}
+
+pub fn rocm_flash_attn_decode(q: &Tensor, k: &Tensor, v: &Tensor, scale: f32) -> Result<Tensor> {
+    let q = q.contiguous()?;
+    q.apply_op3_no_bwd(k, v, &RocmFlashDecode { scale })
+}

hanzo-rocm-kernels/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "hanzo-rocm-kernels"
-version = "0.11.34"
+version = "0.11.35"
 license = "MIT OR Apache-2.0"
 edition = "2021"
 

hanzo-rocm-kernels/src/kernels/flash.hip

@@ -217,6 +217,107 @@ extern "C" __global__ void flash_attn_f16(
     flash_fwd_core<__half>(B, Hq, Hkv, Lq, Lk, scale, causal, q, k, v, o);
 }
 
+// ---------------------------------------------------------------------------------------------
+// Flash-attention DECODE: single query per head (Lq==1) attending the whole KV cache. Fuses the
+// eager path's repeat_kv (cat -> copy2d) + QKᵀ bmm + f32-softmax roundtrip + P@V bmm + the per-layer
+// `.contiguous()` copy of the whole active cache into ONE dispatch. One block per (batch, q_head),
+// 128 threads (= head dim). GQA is implicit (kv_head = q_head / (Hq/Hkv)); the KV cache is read IN
+// PLACE via (batch,head,seq) strides with the innermost head-dim contiguous, so no cache copy. All
+// accumulation is f32: scores, the online (running max/sum) softmax, and the P@V accumulator, so the
+// output matches the eager naive_sdpa within f16 rounding. A decode query sees only past keys, so the
+// full cache is unmasked. Keys are streamed in FD_TILE tiles staged coalesced to LDS.
+#define FD_DH 128
+#define FD_TILE 128
+template <typename T>
+__device__ void flash_decode_core(
+    int B, int Hq, int Hkv, int Lk, float scale,
+    const T* __restrict__ q,
+    const T* __restrict__ k, long long kb_s, long long kh_s, long long kl_s,
+    const T* __restrict__ v, long long vb_s, long long vh_s, long long vl_s,
+    T* __restrict__ o)
+{
+    const int gqa = Hq / Hkv;
+    const int qh = blockIdx.x;          // query head
+    const int b  = blockIdx.y;          // batch
+    const int kvh = qh / gqa;           // GQA-mapped kv head (no repeat_kv)
+    const int t = threadIdx.x;          // 0..127, doubles as key index (scores) then dim index (output)
+
+    __shared__ _Float16 sQ[FD_DH];
+    __shared__ _Float16 sK[FD_TILE * FD_DH];
+    __shared__ float sP[FD_TILE];       // tile probabilities
+    __shared__ float sred[FD_TILE];     // block-reduction scratch
+
+    sQ[t] = (_Float16) fa_ld(q, ((size_t)(b * Hq + qh)) * FD_DH + t);
+
+    const size_t kbase = (size_t)b * (size_t)kb_s + (size_t)kvh * (size_t)kh_s;
+    const size_t vbase = (size_t)b * (size_t)vb_s + (size_t)kvh * (size_t)vh_s;
+
+    float m_run = -1e30f, l_run = 0.0f, o_acc = 0.0f;   // running max, running sum, this thread's O[dim=t]
+    __syncthreads();
+
+    for (int k0 = 0; k0 < Lk; k0 += FD_TILE) {
+        const int tile = min(FD_TILE, Lk - k0);
+        // Coalesced tile load: sK[r][c] = K[batch, kv_head, k0+r, c]  (c = head dim, contiguous).
+        for (int i = t; i < FD_TILE * FD_DH; i += FD_TILE) {
+            const int r = i / FD_DH, c = i % FD_DH;
+            sK[i] = (r < tile) ? (_Float16) fa_ld(k, kbase + (size_t)(k0 + r) * (size_t)kl_s + c) : (_Float16)0;
+        }
+        __syncthreads();
+
+        // score[key=t] = scale * (Q . K[t])
+        float score = -1e30f;
+        if (t < tile) {
+            float s = 0.0f;
+            #pragma unroll
+            for (int d = 0; d < FD_DH; ++d) s += (float)sQ[d] * (float)sK[t * FD_DH + d];
+            score = s * scale;
+        }
+        // block max
+        sred[t] = score; __syncthreads();
+        for (int s = FD_TILE / 2; s > 0; s >>= 1) { if (t < s) sred[t] = fmaxf(sred[t], sred[t + s]); __syncthreads(); }
+        const float tmax = sred[0]; __syncthreads();
+
+        const float m_new = fmaxf(m_run, tmax);
+        const float corr = __expf(m_run - m_new);
+        const float p = (t < tile) ? __expf(score - m_new) : 0.0f;
+        sP[t] = p;
+        // block sum of p
+        sred[t] = p; __syncthreads();
+        for (int s = FD_TILE / 2; s > 0; s >>= 1) { if (t < s) sred[t] += sred[t + s]; __syncthreads(); }
+        l_run = l_run * corr + sred[0];
+        m_run = m_new;
+        __syncthreads();   // sP fully written + sred consumed before the O accumulation / next tile
+
+        // O[dim=t] = O*corr + sum_j p[j] * V[k0+j][t]   (V read coalesced across threads at fixed j)
+        float acc = o_acc * corr;
+        for (int j = 0; j < tile; ++j)
+            acc += sP[j] * (float) fa_ld(v, vbase + (size_t)(k0 + j) * (size_t)vl_s + t);
+        o_acc = acc;
+        __syncthreads();   // before the next tile overwrites sK/sP
+    }
+
+    const float out = (l_run > 0.0f) ? o_acc / l_run : 0.0f;
+    fa_st(o, ((size_t)(b * Hq + qh)) * FD_DH + t, out);
+}
+
+extern "C" __global__ void flash_decode_f16(
+    int B, int Hq, int Hkv, int Lk, float scale,
+    const __half* q,
+    const __half* k, long long kb_s, long long kh_s, long long kl_s,
+    const __half* v, long long vb_s, long long vh_s, long long vl_s,
+    __half* o) {
+    flash_decode_core<__half>(B, Hq, Hkv, Lk, scale, q, k, kb_s, kh_s, kl_s, v, vb_s, vh_s, vl_s, o);
+}
+
+extern "C" __global__ void flash_decode_bf16(
+    int B, int Hq, int Hkv, int Lk, float scale,
+    const __hip_bfloat16* q,
+    const __hip_bfloat16* k, long long kb_s, long long kh_s, long long kl_s,
+    const __hip_bfloat16* v, long long vb_s, long long vh_s, long long vl_s,
+    __hip_bfloat16* o) {
+    flash_decode_core<__hip_bfloat16>(B, Hq, Hkv, Lk, scale, q, k, kb_s, kh_s, kl_s, v, vb_s, vh_s, vl_s, o);
+}
+
 extern "C" __global__ void flash_attn_bf16(
     int B, int Hq, int Hkv, int Lq, int Lk, float scale, int causal,
     const __hip_bfloat16* q, const __hip_bfloat16* k, const __hip_bfloat16* v, __hip_bfloat16* o) {