AstrAI/csrc/kernels/attn_layout_policies.cuh

#pragma once
#include <cuda_bf16.h>
#include "attn_common.h"

// ============================================================================
// Attention layout policies keep Q scheduling independent from K/V storage.
// DenseQSchedule / PackedQSchedule map blocks to Q tiles; ContigKV / PagedKV
// resolve logical K/V positions to physical addresses. This lets the shared
// kernels compose Q layout and K/V storage without coupling the two concerns.
//
//   ContigKV:  K/V are dense [batch, kv_head, kv_len, head_dim] tensors.
//              Params fields used: k, v, kv_stride_*, kv_len, q_len,
//              q_b_stride, causal_offset.
//   PagedKV:   K/V live in a flat pool [size, kv_head, head_dim] indexed via
//              req_to_token.  Params fields used: k_cache, v_cache,
//              req_to_token, req_pool_indices, kv_indptr, qo_indptr,
//              max_context_len, q_l_stride.
//
// Addressing state that is constant across a whole kernel invocation for one
// (batch, kv_head) pair is captured once by make_ctx<HEAD_DIM>() and passed
// to kv_addr, so the load loops never redo the hoistable base computation
// (e.g. the req_pool_indices global read) element-by-element.
// ============================================================================

#define HOST_FORCEINLINE static __host__ __forceinline__
#define DEVICE_FORCEINLINE static __device__ __forceinline__
#define HOST_DEV_FORCEINLINE static __host__ __device__ __forceinline__

using bf16 = __nv_bfloat16;

// ============================================================================
// Q scheduling policies
//
// Map CUDA blocks to request-local Q tiles independently of K/V storage.
// Dense tensors encode the request in blockIdx.z; packed ragged tensors use
// a compact precomputed work map indexed by blockIdx.x.
// ============================================================================

struct DenseQSchedule {
    HOST_FORCEINLINE int host_q_blocks(
        const AttentionParams<bf16>& p, int rows) {
        return (p.q_len + rows - 1) / rows;
    }

    HOST_FORCEINLINE int host_grid_batch(
        const AttentionParams<bf16>& p) {
        return p.batch;
    }

    DEVICE_FORCEINLINE void map_block(
        const AttentionParams<bf16>&, int& batch, int& q_tile) {
        batch = blockIdx.z;
        q_tile = blockIdx.x;
    }

    DEVICE_FORCEINLINE int q_len(
        const AttentionParams<bf16>& p, int) {
        return p.q_len;
    }

    DEVICE_FORCEINLINE int q_base(
        const AttentionParams<bf16>& p, int batch, int q_head) {
        return batch * p.q_b_stride + q_head * p.q_h_stride;
    }
};

struct PackedQSchedule {
    HOST_FORCEINLINE int host_q_blocks(
        const AttentionParams<bf16>& p, int) {
        return p.num_q_tiles;
    }

    HOST_FORCEINLINE int host_grid_batch(
        const AttentionParams<bf16>&) {
        return 1;
    }

    DEVICE_FORCEINLINE void map_block(
        const AttentionParams<bf16>& p, int& batch, int& q_tile) {
        batch = p.q_tile_to_batch[blockIdx.x];
        q_tile = p.q_tile_to_index[blockIdx.x];
    }

    DEVICE_FORCEINLINE int q_len(
        const AttentionParams<bf16>& p, int batch) {
        return p.qo_indptr[batch + 1] - p.qo_indptr[batch];
    }

    DEVICE_FORCEINLINE int q_base(
        const AttentionParams<bf16>& p, int batch, int q_head) {
        return p.qo_indptr[batch] * p.q_l_stride + q_head * p.q_h_stride;
    }
};

// Hoisted per-(batch, kv_head) addressing context.
struct KVContext {
    int kv_base;          // contig: batch*kv_b_stride + kv_head*kv_h_stride
    int req_idx;          // paged: req_pool_indices[batch]
    int64_t rtt_stride;   // paged: max_context_len
    int64_t pool_stride;  // paged: kv_head * HEAD_DIM
    int64_t head_off;     // paged: kv_head * HEAD_DIM
};

// Per-element K/V global addresses for one (kc, d) position of a K/V tile.
// The pointers are ALWAYS the computed addresses (never nullptr) — callers
// gate on `valid` (cp.async src_size=0, or a guarded scalar deref).  `valid`
// starts as "within the request's seq_len"; the paged policy further degrades
// it when req_to_token maps the position to a negative slot (empty padding).
// This matches the original hand-rolled load loops, where the address was
// always formed and the predicate decided whether anything was read.
struct KVAddr {
    const void* k;
    const void* v;
    bool valid;
};

// ---- Contiguous K/V ----
struct ContigKV {
    static constexpr bool kPaged = false;

    HOST_FORCEINLINE int host_kv_len(const AttentionParams<bf16>& p) {
        return p.kv_len;
    }

    // decode: same offset (q_len == 1, so there is no row stride component)
    DEVICE_FORCEINLINE int q_decode_base(
        const AttentionParams<bf16>& p, int batch, int q_head) {
        return batch * p.q_b_stride + q_head * p.q_h_stride;
    }

    DEVICE_FORCEINLINE int kv_len(const AttentionParams<bf16>& p, int) {
        return p.kv_len;
    }
    DEVICE_FORCEINLINE int causal_offset(
        const AttentionParams<bf16>& p, int, int) {
        return p.causal_offset;
    }
    // decode: exclusive bound of the single query's attend range
    DEVICE_FORCEINLINE int decode_attend_len(const AttentionParams<bf16>& p, int) {
        return (p.kv_len < p.causal_offset + 1) ? p.kv_len : (p.causal_offset + 1);
    }

    template <int HEAD_DIM>
    DEVICE_FORCEINLINE KVContext make_ctx(
        const AttentionParams<bf16>& p, int batch, int kv_head) {
        KVContext c = {};
        c.kv_base = batch * p.kv_b_stride + kv_head * p.kv_h_stride;
        return c;
    }
    DEVICE_FORCEINLINE int resolve_token(
        const AttentionParams<bf16>& p, const KVContext& c, int kc, bool valid) {
        return valid ? kc : -1;
    }
    DEVICE_FORCEINLINE KVAddr kv_addr_from_token(
        const AttentionParams<bf16>& p, const KVContext& c, int token, int d) {
        const bool valid = token >= 0;
        const int safe_token = valid ? token : 0;
        const int64_t gmem_off = (int64_t)c.kv_base
            + (int64_t)safe_token * p.kv_l_stride
            + (int64_t)d * p.kv_d_stride;
        return {&p.k_ptr[gmem_off], &p.v_ptr[gmem_off], valid};
    }
};

// ---- Paged (SGLang-style flat pool) K/V ----
struct PagedKV {
    static constexpr bool kPaged = true;

    HOST_FORCEINLINE int host_kv_len(const AttentionParams<bf16>& p) {
        return p.max_context_len;
    }

    // decode: Q is [batch, q_head, head_dim], so batch is the outer row
    DEVICE_FORCEINLINE int q_decode_base(
        const AttentionParams<bf16>& p, int batch, int q_head) {
        return batch * p.q_l_stride + q_head * p.q_h_stride;
    }

    DEVICE_FORCEINLINE int kv_len(const AttentionParams<bf16>& p, int batch) {
        return p.kv_indptr[batch + 1] - p.kv_indptr[batch];
    }
    DEVICE_FORCEINLINE int causal_offset(
        const AttentionParams<bf16>& p, int batch, int q_len) {
        return kv_len(p, batch) - q_len;
    }
    // decode: the query is the last token, so [0, seq_len) IS its causal range
    DEVICE_FORCEINLINE int decode_attend_len(const AttentionParams<bf16>& p, int batch) {
        return kv_len(p, batch);
    }

    template <int HEAD_DIM>
    DEVICE_FORCEINLINE KVContext make_ctx(
        const AttentionParams<bf16>& p, int batch, int kv_head) {
        KVContext c = {};
        c.req_idx = p.req_pool_indices[batch];
        c.rtt_stride = (int64_t)p.max_context_len;
        c.pool_stride = (int64_t)p.kv_head * HEAD_DIM;
        c.head_off = (int64_t)kv_head * HEAD_DIM;
        return c;
    }
    DEVICE_FORCEINLINE int resolve_token(
        const AttentionParams<bf16>& p, const KVContext& c, int kc, bool valid) {
        return valid ? p.req_to_token[c.req_idx * c.rtt_stride + kc] : -1;
    }
    DEVICE_FORCEINLINE KVAddr kv_addr_from_token(
        const AttentionParams<bf16>& p, const KVContext& c, int slot, int d) {
        const bool valid = slot >= 0;
        const int safe_slot = valid ? slot : 0;
        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};
    }
};