refactor: standardize packed 3d inference

- keep training attention on dense 4d tensors
- use packed 3d tensors with KV cache for inference
- extend CUDA rotary embedding to packed 3d inputs
- adapt torch, CUDA and FlashAttention backend dispatch

Benchmark: NVIDIA L20, BF16, 1B model, paged KV cache, CUDA Graph, prompt 512, generation 128 (median of 3 alternating runs)
- batch 1: 234.5 -> 242.6 tok/s (1.034x, +3.4%)
- batch 8: 1243.1 -> 1286.6 tok/s (1.035x, +3.5%)

astrai/extension/backend/attention.py

@@ -253,28 +253,6 @@ def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
     )
 
 
-def _write_and_gather_kv(
-    kv_cache: "KVCache",
-    k: Tensor,
-    v: Tensor,
-    layer_id: int,
-    q: Tensor,
-    attn_mask: Optional[Tensor],
-) -> tuple[Tensor, Tensor]:
-    kv_cache.k_buffer[layer_id, kv_cache.out_cache_loc] = k
-    kv_cache.v_buffer[layer_id, kv_cache.out_cache_loc] = v
-    max_len = kv_cache.max_len
-    indices = kv_cache.req_to_token[kv_cache.req_pool_indices, :max_len]
-    if q.size(1) == 1 and attn_mask is not None and attn_mask.dim() == 4:
-        pos_mask = attn_mask[:, 0, 0]
-    else:
-        pos_mask = (
-            torch.arange(max_len, device=q.device)[None, :] < kv_cache.seq_lens[:, None]
-        )
-    indices = torch.where(pos_mask, indices, torch.zeros_like(indices))
-    return kv_cache.k_buffer[layer_id, indices], kv_cache.v_buffer[layer_id, indices]
-
-
 def attention(
     q: Tensor,
     k: Tensor,
@@ -565,10 +543,6 @@ class CudaBackend(AttentionBackend):
         if kv_cache is None:
             raise RuntimeError("CudaBackend does not support training (kv_cache=None)")
 
-        loc = kv_cache.out_cache_loc
-        kv_cache.k_buffer[layer_id, loc] = k
-        kv_cache.v_buffer[layer_id, loc] = v
-
         kv_indptr = kv_cache.kv_indptr
 
         out = attn_paged_decode(
@@ -578,6 +552,8 @@ class CudaBackend(AttentionBackend):
             kv_cache.req_to_token,
             kv_cache.req_pool_indices,
             kv_indptr,
+            new_k=k,
+            new_v=v,
             is_causal=True,
             o_part_buf=kv_cache.decode_o_part,
             ml_part_buf=kv_cache.decode_ml_part,

astrai/extension/ops/attention.py

@@ -97,6 +97,8 @@ def attn_paged_decode(
     req_to_token: torch.Tensor,
     req_pool_indices: torch.Tensor,
     kv_indptr: torch.Tensor,
+    new_k: Optional[torch.Tensor] = None,
+    new_v: Optional[torch.Tensor] = None,
     mask: Optional[torch.Tensor] = None,
     is_causal: bool = False,
     o_part_buf: Optional[torch.Tensor] = None,
@@ -116,6 +118,8 @@ def attn_paged_decode(
         req_to_token: [num_reqs, max_context_len] (int32) — token -> slot
         req_pool_indices: [batch] (int32) — rows into req_to_token
         kv_indptr: [batch+1] (int32) — prefix sum of per-request seq_lens
+        new_k: current-token K to append, [batch, n_kv_heads, head_dim]
+        new_v: current-token V to append, same shape as new_k
         mask: 2D [batch, max_context_len] (bool, True=keep) or None
         is_causal: apply causal mask
         o_part_buf: pre-allocated split-KV o partial buffer (workflow bypass)
@@ -134,6 +138,8 @@ def attn_paged_decode(
         req_to_token,
         req_pool_indices,
         kv_indptr,
+        new_k=new_k,
+        new_v=new_v,
         mask=mask,
         causal_offset=causal_offset,
         o_part_buf=o_part_buf,

astrai/extension/ops/fp8.py

@@ -27,17 +27,19 @@ def _mod():
 def fp8_mm(
     a: torch.Tensor, b: torch.Tensor, sx: torch.Tensor, sw: torch.Tensor
 ) -> torch.Tensor:
-    """FP8 e4m3 GEMM: a[M,K] x b[N,K] -> bf16[M,N] (pre-scaled inputs)."""
+    """BF16 inputs, fused FP8 GEMM with FP32 accumulation and BF16 output."""
 
 
 @fp8_mm.register_fake
 def _fp8_mm_fake(a, b, sx, sw):
-    return torch.empty((a.size(0), b.size(1)), device=a.device, dtype=torch.bfloat16)
+    return torch.empty((a.size(0), b.size(0)), device=a.device, dtype=torch.bfloat16)
 
 
 @fp8_mm.register_kernel("cuda")
 def _fp8_mm_cuda(a, b, sx, sw):
-    return _mod().fp8_mm(a, b)
+    if not (a.dtype == torch.bfloat16 and b.dtype == torch.bfloat16):
+        raise TypeError(f"bf16 GEMM requires bf16 inputs, got {a.dtype}/{b.dtype}")
+    return _mod().fp8_mm(a, b, sx, sw)
 
 
 @fp8_mm.register_kernel("cpu")
@@ -45,11 +47,63 @@ def _fp8_mm_cpu(a, b, sx, sw):
     return torch.mm(a.float(), b.float().t()).to(torch.bfloat16)
 
 
-def linear_forward_scaled(x, w, bias, sx, sw, sx_inv, sw_inv, amax_x, amax_w):
-    """Quantize x/w with per-tensor scales + cuBLASLt GEMM + bias -> bf16.
+@custom_op("custom::fp8_mm_prequant", mutates_args=())
+def fp8_mm_prequant(
+    a: torch.Tensor, b: torch.Tensor, scale: torch.Tensor
+) -> torch.Tensor:
+    """Pre-quantized FP8 inputs, fused FP8 GEMM, FP32 accumulation, BF16 out."""
 
-    x/w: [..., K] / [N, K] bf16; sx/sw: f32 scale tensors (device scalars);
-    sx_inv/sw_inv: 1/scale; amax_x/amax_w: f32 buffers receiving max-abs.
+
+@fp8_mm_prequant.register_fake
+def _fp8_mm_prequant_fake(a, b, scale):
+    return torch.empty((a.size(0), b.size(0)), device=a.device, dtype=torch.bfloat16)
+
+
+@fp8_mm_prequant.register_kernel("cuda")
+def _fp8_mm_prequant_cuda(a, b, scale):
+    if not (a.dtype == torch.float8_e4m3fn and b.dtype == torch.float8_e4m3fn):
+        raise TypeError(
+            f"pre-quantized FP8 GEMM requires fp8 inputs, got {a.dtype}/{b.dtype}"
+        )
+    return _mod().fp8_mm_prequant(a, b, scale)
+
+
+@fp8_mm_prequant.register_kernel("cpu")
+def _fp8_mm_prequant_cpu(a, b, scale):
+    return (a.float() @ b.float().t() * scale).to(torch.bfloat16)
+
+
+@custom_op("custom::fp8_mm_prequant_fp8", mutates_args=())
+def fp8_mm_prequant_fp8(
+    a: torch.Tensor, b: torch.Tensor, scale: torch.Tensor, out_scale: torch.Tensor
+) -> torch.Tensor:
+    """FP8 inputs and FP8 output: fused FP8 GEMM with FP32 accumulation."""
+
+
+@fp8_mm_prequant_fp8.register_fake
+def _fp8_mm_prequant_fp8_fake(a, b, scale, out_scale):
+    return torch.empty((a.size(0), b.size(0)), device=a.device, dtype=a.dtype)
+
+
+@fp8_mm_prequant_fp8.register_kernel("cuda")
+def _fp8_mm_prequant_fp8_cuda(a, b, scale, out_scale):
+    if not (a.dtype == torch.float8_e4m3fn and b.dtype == torch.float8_e4m3fn):
+        raise TypeError(
+            f"pre-quantized FP8 GEMM requires fp8 inputs, got {a.dtype}/{b.dtype}"
+        )
+    return _mod().fp8_mm_prequant_fp8(a, b, scale, out_scale)
+
+
+@fp8_mm_prequant_fp8.register_kernel("cpu")
+def _fp8_mm_prequant_fp8_cpu(a, b, scale, out_scale):
+    return (a.float() @ b.float().t() * scale * out_scale).to(torch.float8_e4m3fn)
+
+
+def linear_forward_scaled(x, w, bias, sx, sw, sx_inv, sw_inv, amax_x, amax_w):
+    """Quantize BF16 inputs to FP8, accumulate in FP32, and return BF16.
+
+    x/w: [..., K] / [N, K] bf16; sx/sw and their inverses control the fused
+    E4M3 conversion; amax_x/amax_w receive the input max-abs values.
     """
     if not (x.dtype == torch.bfloat16 and w.dtype == torch.bfloat16):
         raise TypeError(f"fp8 forward requires bf16 inputs, got {x.dtype}/{w.dtype}")

csrc/CMakeLists.txt

@@ -61,9 +61,6 @@ foreach(name ${KERNELS})
         "${PYTHON_INCLUDE_DIR}")
 
     target_link_libraries(${name} PRIVATE ${TORCH_LIBS})
-    if(${name} STREQUAL "fp8_mm")
-        target_link_libraries(${name} PRIVATE CUDA::cublasLt)
-    endif()
     target_link_options(${name} PRIVATE "-Wl,-rpath,${TORCH_LIB_DIR}")
 
     target_compile_options(${name} PRIVATE

csrc/kernels/attn_common.h

@@ -36,6 +36,8 @@ struct AttentionParams {
     const T* __restrict__ q_ptr;
     const T* __restrict__ k_ptr;
     const T* __restrict__ v_ptr;
+    const T* __restrict__ new_k_ptr;
+    const T* __restrict__ new_v_ptr;
     T* __restrict__ o_ptr;
     const bool* __restrict__ mask;
 
@@ -50,6 +52,9 @@ struct AttentionParams {
     int kv_l_stride;
     int kv_d_stride;
 
+    int new_kv_b_stride;
+    int new_kv_h_stride;
+
     int mask_b_stride;
     int mask_h_stride;
     int mask_l_stride;

csrc/kernels/attn_decode_split_kv.cuh

@@ -57,8 +57,8 @@ __global__ void attn_decode_split_kv_kernel(AttentionParams<bf16> p) {
             int s = i / p.head_dim;
             int d_dim = i % p.head_dim;
             int kc = chunk_start + s;
-            int token = KV::resolve_token(p, kctx, kc, true);
-            KVAddr a = KV::kv_addr_from_token(p, kctx, token, d_dim);
+            KVAddr a = KV::template decode_addr<1>(
+                p, kctx, batch, kv_head, kc, d_dim, true, true);
             k_smem[i] = a.valid ? *reinterpret_cast<const bf16*>(a.k) : (bf16)0.f;
             v_smem[i] = a.valid ? *reinterpret_cast<const bf16*>(a.v) : (bf16)0.f;
         }

csrc/kernels/attn_decode_split_kv_mma.cuh

@@ -73,8 +73,10 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
             int r = i / Traits::HEAD_DIM, d = i % Traits::HEAD_DIM;
             int kc = kv0 + r;
             bool valid = kc < seq_len;
-            int token = KV::resolve_token(p, kctx, kc, valid);
-            KVAddr a = KV::kv_addr_from_token(p, kctx, token, d);
+            // All GQA passes consume new K/V directly. Only the first pass
+            // persists it, so no cross-block synchronization is required.
+            KVAddr a = KV::template decode_addr<Traits::VEC>(
+                p, kctx, batch, kv_head, kc, d, valid, pass == 0);
             int off = r * Traits::LD + swiz_col(d, r, Traits::SWIZ_MASK);
             cp_async_16_pred(&dK[off], a.k, a.valid);
             cp_async_16_pred(&dV[off], a.v, a.valid);

csrc/kernels/attn_entry_utils.cuh

@@ -130,6 +130,8 @@ inline void attn_pack_params(
     p.q_ptr = (const T*)q.data_ptr();
     p.k_ptr = (const T*)k.data_ptr();
     p.v_ptr = (const T*)v.data_ptr();
+    p.new_k_ptr = nullptr;
+    p.new_v_ptr = nullptr;
     p.o_ptr = nullptr;
     p.o_part = nullptr;
     p.ml_part = nullptr;
@@ -148,6 +150,8 @@ inline void attn_pack_paged_decode_params(
     torch::Tensor req_to_token,
     torch::Tensor req_pool_indices,
     torch::Tensor kv_indptr,
+    const c10::optional<torch::Tensor>& new_k,
+    const c10::optional<torch::Tensor>& new_v,
     c10::optional<torch::Tensor> mask,
     int64_t causal_offset,
     double scale,
@@ -191,6 +195,33 @@ inline void attn_pack_paged_decode_params(
     p.qo_indptr = nullptr;
     p.max_context_len = (int)req_to_token.size(1);
 
+    TORCH_CHECK(new_k.has_value() == new_v.has_value(),
+                "new_k and new_v must be provided together");
+    if (new_k.has_value()) {
+        auto nk = new_k.value();
+        auto nv = new_v.value();
+        TORCH_CHECK(nk.is_cuda() && nv.is_cuda(), "new K/V must be CUDA tensors");
+        TORCH_CHECK(nk.dtype() == torch::kBFloat16 && nv.dtype() == torch::kBFloat16,
+                    "new K/V must be bf16");
+        TORCH_CHECK(nk.dim() == 3 && nv.dim() == 3,
+                    "new K/V must be 3D [batch, kv_head, head_dim]");
+        TORCH_CHECK(nk.sizes() == nv.sizes(), "new K and V must have identical shapes");
+        TORCH_CHECK(nk.strides() == nv.strides(),
+                    "new K and V must have identical strides");
+        TORCH_CHECK(nk.size(0) == p.batch && nk.size(1) == p.kv_head
+                    && nk.size(2) == p.head_dim, "new K/V shape mismatch");
+        TORCH_CHECK(nk.stride(2) == 1 && nv.stride(2) == 1,
+                    "new K/V head_dim must be contiguous");
+        p.new_k_ptr = (const T*)nk.data_ptr();
+        p.new_v_ptr = (const T*)nv.data_ptr();
+        p.new_kv_b_stride = (int)nk.stride(0);
+        p.new_kv_h_stride = (int)nk.stride(1);
+    } else {
+        p.new_k_ptr = nullptr;
+        p.new_v_ptr = nullptr;
+        p.new_kv_b_stride = p.new_kv_h_stride = 0;
+    }
+
     p.causal_offset = (int)causal_offset;
     p.use_mask = (mask.has_value() && mask.value().defined()) ? 1 : 0;
     p.scale = (scale > 0.0) ? (float)scale : 1.0f / sqrtf((float)p.head_dim);
@@ -279,6 +310,8 @@ inline void attn_pack_paged_prefill_params(
 
     p.k_ptr = (const T*)k_cache.data_ptr();
     p.v_ptr = (const T*)v_cache.data_ptr();
+    p.new_k_ptr = nullptr;
+    p.new_v_ptr = nullptr;
     p.q_ptr = (const T*)q.data_ptr();
     p.req_to_token = req_to_token.data_ptr<int>();
     p.req_pool_indices = req_pool_indices.data_ptr<int>();

csrc/kernels/attn_layout_policies.cuh

@@ -160,6 +160,14 @@ struct ContigKV {
             + (int64_t)d * p.kv_d_stride;
         return {&p.k_ptr[gmem_off], &p.v_ptr[gmem_off], valid};
     }
+
+    template <int VEC>
+    DEVICE_FORCEINLINE KVAddr decode_addr(
+        const AttentionParams<bf16>& p, const KVContext& c,
+        int, int, int kc, int d, bool valid, bool) {
+        int token = resolve_token(p, c, kc, valid);
+        return kv_addr_from_token(p, c, token, d);
+    }
 };
 
 // ---- Paged (SGLang-style flat pool) K/V ----
@@ -209,4 +217,39 @@ struct PagedKV {
         const int64_t gmem_off = (int64_t)safe_slot * c.pool_stride + c.head_off + d;
         return {&p.k_ptr[gmem_off], &p.v_ptr[gmem_off], valid};
     }
+
+    DEVICE_FORCEINLINE KVAddr new_kv_addr(
+        const AttentionParams<bf16>& p, int batch, int kv_head, int d) {
+        const int64_t off = (int64_t)batch * p.new_kv_b_stride
+            + (int64_t)kv_head * p.new_kv_h_stride + d;
+        return {&p.new_k_ptr[off], &p.new_v_ptr[off], true};
+    }
+
+    DEVICE_FORCEINLINE void store_new_kv(
+        const AttentionParams<bf16>& p, const KVContext& c,
+        int seq_len, int d, const KVAddr& src) {
+        int slot = resolve_token(p, c, seq_len - 1, true);
+        const int64_t off = (int64_t)slot * c.pool_stride + c.head_off + d;
+        const_cast<bf16*>(p.k_ptr)[off] = *reinterpret_cast<const bf16*>(src.k);
+        const_cast<bf16*>(p.v_ptr)[off] = *reinterpret_cast<const bf16*>(src.v);
+    }
+
+    template <int VEC>
+    DEVICE_FORCEINLINE KVAddr decode_addr(
+        const AttentionParams<bf16>& p, const KVContext& c,
+        int batch, int kv_head, int kc, int d, bool valid, bool persist) {
+        if (p.new_k_ptr && valid && kc == kv_len(p, batch) - 1) {
+            KVAddr src = new_kv_addr(p, batch, kv_head, d);
+            if (persist) {
+                #pragma unroll
+                for (int j = 0; j < VEC; j++) {
+                    KVAddr value = new_kv_addr(p, batch, kv_head, d + j);
+                    store_new_kv(p, c, kc + 1, d + j, value);
+                }
+            }
+            return src;
+        }
+        int token = resolve_token(p, c, kc, valid);
+        return kv_addr_from_token(p, c, token, d);
+    }
 };

csrc/kernels/attn_paged_decode.cu

@@ -8,6 +8,8 @@ torch::Tensor attn_paged_decode(
     torch::Tensor req_to_token,
     torch::Tensor req_pool_indices,
     torch::Tensor kv_indptr,
+    c10::optional<torch::Tensor> new_k,
+    c10::optional<torch::Tensor> new_v,
     c10::optional<torch::Tensor> mask,
     int64_t causal_offset,
     double scale,
@@ -21,6 +23,7 @@ torch::Tensor attn_paged_decode(
     AttentionParams<bf16> p;
     attn_pack_paged_decode_params(q, k_cache, v_cache,
                                    req_to_token, req_pool_indices, kv_indptr,
+                                   new_k, new_v,
                                    mask, causal_offset, scale, p);
 
     torch::Tensor O;
@@ -71,6 +74,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         py::arg("req_to_token"),
         py::arg("req_pool_indices"),
         py::arg("kv_indptr"),
+        py::arg("new_k") = py::none(),
+        py::arg("new_v") = py::none(),
         py::arg("mask") = py::none(),
         py::arg("causal_offset") = -1,
         py::arg("scale") = 0.0,

csrc/tests/fp8_mma_test.cu

@@ -0,0 +1,155 @@
+/*
+Single-kernel BF16 -> FP8 MMA -> BF16 demo for Ada (sm_89).
+
+nvcc -I csrc -arch=sm_89 -std=c++17 -O3 --use_fast_math \
+    --ptxas-options=-O3,-v csrc/tests/fp8_mma_test.cu -o fp8_mma_test \
+    && ./fp8_mma_test
+*/
+
+#include "test_utils.cuh"
+
+#include <cuda_fp8.h>
+
+#include <algorithm>
+#include <vector>
+
+constexpr int M = 16;
+constexpr int N = 8;
+constexpr int K = 32;
+
+__device__ __forceinline__ unsigned pack_fp8x4(float x0, float x1, float x2,
+                                                float x3) {
+    __nv_fp8_e4m3 q0(x0);
+    __nv_fp8_e4m3 q1(x1);
+    __nv_fp8_e4m3 q2(x2);
+    __nv_fp8_e4m3 q3(x3);
+    return static_cast<unsigned>(q0.__x) |
+           (static_cast<unsigned>(q1.__x) << 8) |
+           (static_cast<unsigned>(q2.__x) << 16) |
+           (static_cast<unsigned>(q3.__x) << 24);
+}
+
+__device__ __forceinline__ unsigned load_quantize_fp8x4(
+    const bf16* src, float scale_inv) {
+    return pack_fp8x4(__bfloat162float(src[0]) * scale_inv,
+                      __bfloat162float(src[1]) * scale_inv,
+                      __bfloat162float(src[2]) * scale_inv,
+                      __bfloat162float(src[3]) * scale_inv);
+}
+
+__device__ __forceinline__ void mma_fp8_16832(float d[4],
+                                               const unsigned a[4],
+                                               const unsigned b[2]) {
+    asm volatile(
+        "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
+        "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};"
+        : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3])
+        : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]),
+          "r"(b[0]), "r"(b[1]));
+}
+
+__global__ void fused_bf16_fp8_mma_kernel(
+    const bf16* __restrict__ a, const bf16* __restrict__ b,
+    bf16* __restrict__ out, float scale_a, float scale_b) {
+    const int lane = threadIdx.x;
+    const int group = lane >> 2;
+    const int thread_in_group = lane & 3;
+    const int k0 = thread_in_group * 4;
+
+    // PTX m16n8k32 A fragment: two rows, two 16-column K partitions.
+    unsigned a_frag[4];
+    a_frag[0] = load_quantize_fp8x4(&a[group * K + k0], 1.0f / scale_a);
+    a_frag[1] = load_quantize_fp8x4(&a[(group + 8) * K + k0], 1.0f / scale_a);
+    a_frag[2] = load_quantize_fp8x4(&a[group * K + k0 + 16], 1.0f / scale_a);
+    a_frag[3] = load_quantize_fp8x4(&a[(group + 8) * K + k0 + 16],
+                                    1.0f / scale_a);
+
+    // B is supplied as row-major [N,K], equivalent to the col-major [K,N]
+    // operand required by the MMA instruction.
+    unsigned b_frag[2];
+    b_frag[0] = load_quantize_fp8x4(&b[group * K + k0], 1.0f / scale_b);
+    b_frag[1] = load_quantize_fp8x4(&b[group * K + k0 + 16], 1.0f / scale_b);
+
+    float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f};
+    mma_fp8_16832(acc, a_frag, b_frag);
+
+    const int col = thread_in_group * 2;
+    const float output_scale = scale_a * scale_b;
+    *reinterpret_cast<__nv_bfloat162*>(&out[group * N + col]) =
+        __floats2bfloat162_rn(acc[0] * output_scale,
+                              acc[1] * output_scale);
+    *reinterpret_cast<__nv_bfloat162*>(&out[(group + 8) * N + col]) =
+        __floats2bfloat162_rn(acc[2] * output_scale,
+                              acc[3] * output_scale);
+}
+
+static float quantize_e4m3(float value) {
+    return static_cast<float>(__nv_fp8_e4m3(value));
+}
+
+int main() {
+    srand(0);
+    std::vector<float> a(M * K), b(N * K), reference(M * N, 0.0f);
+    std::vector<bf16> a_bf16(M * K), b_bf16(N * K), output(M * N);
+    for (float& value : a) value = randf() * 4.0f;
+    for (float& value : b) value = randf() * 4.0f;
+    for (int i = 0; i < M * K; ++i) {
+        a_bf16[i] = f2bf(a[i]);
+        a[i] = bf2f(a_bf16[i]);
+    }
+    for (int i = 0; i < N * K; ++i) {
+        b_bf16[i] = f2bf(b[i]);
+        b[i] = bf2f(b_bf16[i]);
+    }
+
+    const float amax = *std::max_element(
+        a.begin(), a.end(), [](float x, float y) { return fabsf(x) < fabsf(y); });
+    const float bmax = *std::max_element(
+        b.begin(), b.end(), [](float x, float y) { return fabsf(x) < fabsf(y); });
+    const float scale_a = fabsf(amax) / 448.0f;
+    const float scale_b = fabsf(bmax) / 448.0f;
+
+    for (int row = 0; row < M; ++row) {
+        for (int col = 0; col < N; ++col) {
+            float sum = 0.0f;
+            for (int k = 0; k < K; ++k) {
+                float qa = quantize_e4m3(a[row * K + k] / scale_a);
+                float qb = quantize_e4m3(b[col * K + k] / scale_b);
+                sum = fmaf(qa, qb, sum);
+            }
+            reference[row * N + col] = sum * scale_a * scale_b;
+        }
+    }
+
+    bf16 *d_a, *d_b, *d_out;
+    CUDA_CHECK(cudaMalloc(&d_a, a_bf16.size() * sizeof(bf16)));
+    CUDA_CHECK(cudaMalloc(&d_b, b_bf16.size() * sizeof(bf16)));
+    CUDA_CHECK(cudaMalloc(&d_out, output.size() * sizeof(bf16)));
+    CUDA_CHECK(cudaMemcpy(d_a, a_bf16.data(), a_bf16.size() * sizeof(bf16),
+                          cudaMemcpyHostToDevice));
+    CUDA_CHECK(cudaMemcpy(d_b, b_bf16.data(), b_bf16.size() * sizeof(bf16),
+                          cudaMemcpyHostToDevice));
+
+    fused_bf16_fp8_mma_kernel<<<1, 32>>>(d_a, d_b, d_out, scale_a, scale_b);
+    CUDA_CHECK(cudaDeviceSynchronize());
+    CUDA_CHECK(cudaMemcpy(output.data(), d_out, output.size() * sizeof(bf16),
+                          cudaMemcpyDeviceToHost));
+
+    float max_abs_error = 0.0f;
+    float max_rel_error = 0.0f;
+    for (int i = 0; i < M * N; ++i) {
+        float error = fabsf(bf2f(output[i]) - reference[i]);
+        max_abs_error = fmaxf(max_abs_error, error);
+        max_rel_error = fmaxf(max_rel_error,
+                              error / fmaxf(fabsf(reference[i]), 1e-4f));
+    }
+    const bool pass = max_abs_error < 0.05f;
+    print_test_header();
+    print_test_row("M=16 N=8 K=32 fused BF16->E4M3 MMA", max_abs_error,
+                   max_rel_error, pass);
+
+    cudaFree(d_a);
+    cudaFree(d_b);
+    cudaFree(d_out);
+    return pass ? 0 : 1;
+}

tests/extension/test_backend_equivalence.py

@@ -7,6 +7,7 @@ seq_lens with padding mask), and end-to-end scheduler.run_batch.
 import torch
 
 from astrai.extension import ATTN_BACKEND, attn_backend
+from astrai.extension.ops.attention import attn_paged_decode
 from astrai.inference.cache import PagePool, TaskCacheManager
 from astrai.inference.runtime.graph import CudaGraphContext
 from astrai.inference.scheduler import InferenceScheduler
@@ -160,6 +161,56 @@ def test_decode_mixed_seq_lens_matches_torch(cuda_model):
     assert diff < 0.05, f"Decode diff (mixed seq_lens): {diff}"
 
 
+@skip_no_kernel
+def test_paged_decode_appends_new_kv_in_kernel():
+    """Fused decode writes current-token K/V to each request's paged slot."""
+    pool = PagePool(
+        n_layers=1,
+        n_kv_heads=1,
+        head_dim=D,
+        max_batch_size=2,
+        max_seq_len=64,
+        device="cuda",
+        dtype=torch.bfloat16,
+        page_size=8,
+        n_tokens=128,
+    )
+    task_cache = _mk_task_cache(pool)
+    ws = _ws(pool)
+    task_cache.task_alloc("t1", list(range(8)))
+    task_cache.task_alloc("t2", list(range(6)))
+    task_cache.task_extend("t1", 8)
+    task_cache.task_extend("t2", 6)
+    kv_cache = task_cache.bind(["t1", "t2"], ws)
+
+    q = torch.randn(2, 2, D, device="cuda", dtype=torch.bfloat16)
+    new_k = torch.randn(2, 1, D, device="cuda", dtype=torch.bfloat16)
+    new_v = torch.randn(2, 1, D, device="cuda", dtype=torch.bfloat16)
+    out = attn_paged_decode(
+        q,
+        kv_cache.k_buffer[0],
+        kv_cache.v_buffer[0],
+        kv_cache.req_to_token,
+        kv_cache.req_pool_indices,
+        kv_cache.kv_indptr,
+        new_k=new_k,
+        new_v=new_v,
+        is_causal=True,
+        o_part_buf=kv_cache.decode_o_part,
+        ml_part_buf=kv_cache.decode_ml_part,
+        out_buf=kv_cache.decode_out,
+    )
+    torch.cuda.synchronize()
+
+    torch.testing.assert_close(
+        kv_cache.k_buffer[0, kv_cache.out_cache_loc], new_k, rtol=0, atol=0
+    )
+    torch.testing.assert_close(
+        kv_cache.v_buffer[0, kv_cache.out_cache_loc], new_v, rtol=0, atol=0
+    )
+    assert torch.isfinite(out).all()
+
+
 @skip_no_kernel
 def test_decode_cuda_graph_replay_is_exact(cuda_model):
     """INT32 cache indices must remain graph-capturable and replay exactly."""

tests/extension/test_fp8_mma.py

@@ -0,0 +1,156 @@
+"""Fused BF16-boundary FP8 MMA kernel tests."""
+
+import pytest
+import torch
+
+from astrai.extension.loader import get_module, is_available
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available()
+    or torch.cuda.get_device_capability() < (8, 9)
+    or not is_available("fp8_mm"),
+    reason="fused FP8 MMA requires a built kernel and compute capability 8.9+",
+)
+
+
+def _scale(tensor):
+    return (tensor.abs().amax().float() / 448.0).clamp_min(1e-12)
+
+
+def _quantize(tensor, scale):
+    return (tensor.float() / scale).to(torch.float8_e4m3fn).float()
+
+
+@pytest.mark.parametrize(
+    ("m", "n", "k"),
+    [(16, 8, 32), (17, 9, 33), (31, 15, 64), (32, 48, 96)],
+)
+def test_fused_fp8_mma_matches_explicit_quantization(m, n, k):
+    torch.manual_seed(m + n + k)
+    a = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
+    b = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
+    scale_a = _scale(a)
+    scale_b = _scale(b)
+
+    out = get_module("fp8_mm").fp8_mm(a, b, scale_a, scale_b)
+    expected = (
+        _quantize(a, scale_a) @ _quantize(b, scale_b).t() * scale_a * scale_b
+    ).to(torch.bfloat16)
+
+    assert out.dtype == torch.bfloat16
+    assert out.shape == (m, n)
+    torch.testing.assert_close(out, expected, atol=0.125, rtol=0.01)
+
+
+def test_fused_fp8_linear_forward_and_backward():
+    torch.manual_seed(7)
+    m, n, k = 19, 13, 37
+    x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
+    weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
+    grad = torch.randn(m, n, device="cuda", dtype=torch.bfloat16)
+    bias = torch.randn(n, device="cuda", dtype=torch.bfloat16)
+    scale_x, scale_w, scale_g = _scale(x), _scale(weight), _scale(grad)
+    amax_x = torch.empty(1, device="cuda", dtype=torch.float32)
+    amax_w = torch.empty(1, device="cuda", dtype=torch.float32)
+    amax_g = torch.empty(1, device="cuda", dtype=torch.float32)
+    module = get_module("fp8_mm")
+
+    out = module.fp8_linear_forward_scaled(
+        x,
+        weight,
+        bias,
+        scale_x,
+        scale_w,
+        scale_x.reciprocal(),
+        scale_w.reciprocal(),
+        amax_x,
+        amax_w,
+    )
+    grad_x, grad_w, grad_b = module.fp8_linear_backward_scaled(
+        grad,
+        x,
+        weight,
+        [1, 1, 1],
+        scale_g,
+        scale_w,
+        scale_x,
+        scale_g.reciprocal(),
+        scale_w.reciprocal(),
+        scale_x.reciprocal(),
+        amax_g,
+    )
+
+    qx = _quantize(x, scale_x)
+    qw = _quantize(weight, scale_w)
+    qg = _quantize(grad, scale_g)
+    expected_out = (qx @ qw.t() * scale_x * scale_w + bias).to(torch.bfloat16)
+    expected_grad_x = (qg @ qw * scale_g * scale_w).to(torch.bfloat16)
+    expected_grad_w = (qg.t() @ qx * scale_g * scale_x).to(torch.bfloat16)
+
+    torch.testing.assert_close(out, expected_out, atol=0.125, rtol=0.01)
+    torch.testing.assert_close(grad_x, expected_grad_x, atol=0.125, rtol=0.01)
+    torch.testing.assert_close(grad_w, expected_grad_w, atol=0.125, rtol=0.01)
+    torch.testing.assert_close(grad_b, grad.sum(0).to(torch.bfloat16))
+    torch.testing.assert_close(amax_x, x.abs().amax().float().reshape(1))
+    torch.testing.assert_close(amax_w, weight.abs().amax().float().reshape(1))
+    torch.testing.assert_close(amax_g, grad.abs().amax().float().reshape(1))
+
+
+def test_fp8_mm_prequant_matches_scaled_mm():
+    torch.manual_seed(11)
+    m, n, k = 512, 4096, 4096
+    a = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
+    weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
+    a8 = a.to(torch.float8_e4m3fn)
+    w8 = weight.to(torch.float8_e4m3fn)
+    scale = torch.tensor([2.5], device="cuda")
+
+    out = get_module("fp8_mm").fp8_mm_prequant(a8, w8, scale)
+
+    # Reference via fp64: FP8 quantization error is dominated by the 3-bit
+    # mantissa, so the tolerance must track the input quantization scale.
+    ref = (a8.float().double() @ w8.float().double().t() * 2.5).to(torch.bfloat16)
+    assert out.dtype == torch.bfloat16
+    assert out.shape == (m, n)
+    torch.testing.assert_close(out, ref, atol=6.0, rtol=0.05)
+
+    # Cross-check against torch's native FP8 GEMM on identical inputs.
+    try:
+        torch._scaled_mm(
+            a8,
+            w8.t(),
+            torch.full((m, 1), 2.5, device="cuda"),
+            torch.ones((1, n), device="cuda"),
+            out_dtype=torch.bfloat16,
+        )
+    except (RuntimeError, NotImplementedError):
+        return
+    torch.testing.assert_close(
+        out,
+        torch._scaled_mm(
+            a8,
+            w8.t(),
+            torch.full((m, 1), 2.5, device="cuda"),
+            torch.ones((1, n), device="cuda"),
+            out_dtype=torch.bfloat16,
+        ),
+        atol=2.0,
+        rtol=0.01,
+    )
+
+
+def test_fp8_mm_prequant_fp8_output():
+    torch.manual_seed(13)
+    m, n, k = 512, 4096, 4096
+    a = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
+    weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
+    a8 = a.to(torch.float8_e4m3fn)
+    w8 = weight.to(torch.float8_e4m3fn)
+    scale = torch.tensor([2.5], device="cuda")
+    out_scale = torch.tensor([0.1], device="cuda")
+
+    out = get_module("fp8_mm").fp8_mm_prequant_fp8(a8, w8, scale, out_scale)
+    assert out.dtype == torch.float8_e4m3fn
+    assert out.shape == (m, n)
+    ref = (a8.float().double() @ w8.float().double().t() * 2.5 * 0.1).to(torch.bfloat16)
+    torch.testing.assert_close(out.float().to(torch.bfloat16), ref, atol=1.0, rtol=0.05)

tests/trainer/test_grpo_e2e.py

@@ -0,0 +1,123 @@
+"""End-to-end integration test for online GRPO rollout."""
+
+import os
+from functools import partial
+
+import pytest
+import torch
+from torch.utils.data import Dataset
+
+from astrai.config import TrainConfig
+from astrai.model.transformer import AutoRegressiveLM
+from astrai.trainer.rollout import BaseRewardModel
+from astrai.trainer.schedule import SchedulerFactory
+from astrai.trainer.trainer import Trainer
+from tests.helpers import CHAT_TEMPLATE
+
+
+class InstructionDataset(Dataset):
+    """Toy instruction/input dataset for online RL rollout.
+
+    Each sample has an ``instruction`` and an optional ``input``; the
+    RolloutGenerator renders both through the tokenizer's chat template
+    so the prompt matches the SFT-trained format.
+    """
+
+    _SAMPLES = [
+        {"instruction": "Hello", "input": ""},
+        {"instruction": "Tell me a story", "input": "about dragons"},
+        {"instruction": "Summarize", "input": "the article"},
+        {"instruction": "Translate", "input": "to French: hi"},
+    ]
+
+    def __len__(self):
+        return len(self._SAMPLES)
+
+    def __getitem__(self, idx):
+        return dict(self._SAMPLES[idx])
+
+
+class LengthRewardModel(BaseRewardModel):
+    """Rewards each response by its (non-pad) token count.
+
+    Gives the group-normalized advantage a non-degenerate signal.
+    """
+
+    def score(self, prompts, responses):
+        B = len(prompts)
+        G = len(responses[0]) if B else 0
+        rewards = torch.zeros(B, G)
+        for i in range(B):
+            for g in range(G):
+                rewards[i, g] = float(len(responses[i][g]))
+        return rewards
+
+
+def instruction_collate_fn(batch):
+    """Stack a list of instruction/input dicts into a batch dict of lists."""
+    return {
+        "instruction": [b["instruction"] for b in batch],
+        "input": [b.get("input", "") for b in batch],
+    }
+
+
+def _model_fn(model_config):
+    return AutoRegressiveLM(model_config).to(dtype=torch.float32)
+
+
+def _optimizer_fn(m):
+    return torch.optim.AdamW(m.parameters(), lr=1e-4)
+
+
+def _scheduler_fn(optim):
+    return SchedulerFactory.create(
+        "cosine", optim, warmup_steps=1, lr_decay_steps=4, min_rate=0.05
+    )
+
+
+@pytest.mark.integration
+def test_online_grpo_end_to_end(base_test_env):
+    """Run one epoch of online GRPO with KV-cache-backed rollout."""
+    test_dir = base_test_env["test_dir"]
+    device = base_test_env["device"]
+    tokenizer = base_test_env["tokenizer"]
+    model_config = base_test_env["transformer_config"]
+
+    tokenizer.set_chat_template(CHAT_TEMPLATE)
+    tokenizer.save_pretrained(test_dir)
+
+    model_fn = partial(_model_fn, model_config)
+    optimizer_fn = _optimizer_fn
+    scheduler_fn = _scheduler_fn
+
+    dataset = InstructionDataset()
+
+    train_config = TrainConfig(
+        strategy="online_grpo",
+        model_fn=model_fn,
+        dataset=dataset,
+        optimizer_fn=optimizer_fn,
+        scheduler_fn=scheduler_fn,
+        ckpt_dir=os.path.join(test_dir, "ckpt"),
+        n_epoch=1,
+        batch_per_device=2,
+        ckpt_interval=100,
+        grad_accum_steps=1,
+        random_seed=42,
+        device_type=device,
+        nprocs=1,
+        parallel_mode="none",
+        extra_kwargs={"clip_eps": 0.2, "kl_coef": 0.01, "group_size": 2},
+        rollout_interval=1,
+        rollout_temperature=1.0,
+        rollout_top_k=0,
+        rollout_top_p=1.0,
+        rollout_max_tokens=4,
+        reward_model_fn=LengthRewardModel,
+        collate_fn=instruction_collate_fn,
+    )
+
+    trainer = Trainer(train_config)
+    trainer.train(param_path=test_dir)
+
+    assert os.path.isdir(os.path.join(test_dir, "ckpt"))