hey! this is a static git repo viewer. commit diffs, file history, older commits and other revisions won't load without javascript.

you are still able to view source files on branch/tag tips, but the experience is degraded.

char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

ynankaniCUDA: XOR swizzle flash attn K,V smem fp16 tiles (#25635)e4b9af0072026-08-31

1#pragma once
2
3#include "common.cuh"
4#include "mma.cuh"
5
6// XOR swizzle for K/V SMEM tiles to avoid bank conflicts without row padding (Turing+ only).
7// Stride must be a multiple of 32 half2 columns, otherwise we keep +4 row padding.
8
9namespace ggml_cuda_fattn_smem_swizzle {
10
11static __host__ __device__ constexpr bool bank_aligned(const int nbatch_2) {
12    return nbatch_2 >= 32 && nbatch_2 % 32 == 0;
13}
14
15static __device__ constexpr bool enabled(const int nbatch_2) {
16#if defined(TURING_MMA_AVAILABLE)
17    return bank_aligned(nbatch_2);
18#else
19    GGML_UNUSED(nbatch_2);
20    return false;
21#endif // defined(TURING_MMA_AVAILABLE)
22}
23
24static __host__ bool enabled(const int nbatch_2, const int cc) {
25#ifdef GGML_USE_HIP
26    GGML_UNUSED(nbatch_2);
27    GGML_UNUSED(cc);
28    return false;
29#else
30    return turing_mma_available(cc) && bank_aligned(nbatch_2);
31#endif // GGML_USE_HIP
32}
33
34static __device__ constexpr int tile_stride(const int nbatch_2) {
35    return enabled(nbatch_2) ? nbatch_2 : nbatch_2 + 4;
36}
37
38static __host__ int tile_stride(const int nbatch_2, const int cc) {
39    return enabled(nbatch_2, cc) ? nbatch_2 : nbatch_2 + 4;
40}
41
42// Swizzled byte offset for tile element (row, col_h2), same map used for writes and reads.
43template<int stride_h2>
44static __device__ __forceinline__ int bytes_rc(const int row, const int col_h2) {
45    static_assert(bank_aligned(stride_h2), "swizzled tile needs a stride that is a multiple of 32");
46    return ((row * stride_h2 + col_h2) * (int) sizeof(half2)) ^ ((row & 7) << 4);
47}
48
49// ldmatrix.x4 via 64-bit generic pointer.
50static __device__ __forceinline__ void ldmatrix_x4(int * xi, const half2 * addr) {
51#if defined(TURING_MMA_AVAILABLE)
52    asm volatile("ldmatrix.sync.aligned.m8n8.x4.b16 {%0, %1, %2, %3}, [%4];"
53        : "=r"(xi[0]), "=r"(xi[1]), "=r"(xi[2]), "=r"(xi[3])
54        : "l"(addr));
55#else
56    GGML_UNUSED_VARS(xi, addr);
57    NO_DEVICE_CODE;
58#endif // defined(TURING_MMA_AVAILABLE)
59}
60
61static __device__ __forceinline__ void ldmatrix_x4_trans(int * xi, const half2 * addr) {
62#if defined(TURING_MMA_AVAILABLE)
63    asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.b16 {%0, %1, %2, %3}, [%4];"
64        : "=r"(xi[0]), "=r"(xi[2]), "=r"(xi[1]), "=r"(xi[3])
65        : "l"(addr));
66#else
67    GGML_UNUSED_VARS(xi, addr);
68    NO_DEVICE_CODE;
69#endif // defined(TURING_MMA_AVAILABLE)
70}
71
72// Per-lane swizzled address for one tile<16, 8, half2> ldmatrix: 16 rows, 4 half2 columns per lane.
73template<int stride_h2>
74static __device__ __forceinline__ const half2 * lane_addr(
75        const half2 * tile_base, const int base_row, const int base_col_h2, const int I, const int J) {
76    static_assert(bank_aligned(stride_h2), "swizzled tile needs a stride that is a multiple of 32");
77    const int lane_row = threadIdx.x % I;
78    const int lane_col = (threadIdx.x / I) * (J / 2);
79    uint32_t byte_off = (uint32_t) ((base_row + lane_row)*stride_h2 + base_col_h2 + lane_col) * (uint32_t) sizeof(half2);
80    byte_off ^= (uint32_t) (((base_row + lane_row) & 7) << 4);
81    return (const half2 *) ((const char *) tile_base + byte_off);
82}
83
84template<int stride_h2, bool swz, typename TileT>
85static __device__ __forceinline__ void load_ldmatrix(
86        TileT & t, const half2 * tile_base, const int base_row, const int base_col_h2) {
87    if constexpr (swz) {
88        static_assert(std::is_same_v<TileT, ggml_cuda_mma::tile<16, 8, half2>>,
89            "the swizzled layout is only supported for tile<16, 8, half2>");
90        ldmatrix_x4((int *) t.x, lane_addr<stride_h2>(tile_base, base_row, base_col_h2, TileT::I, TileT::J));
91    } else {
92        ggml_cuda_mma::load_ldmatrix(t, tile_base + base_row*stride_h2 + base_col_h2, stride_h2);
93    }
94}
95
96template<int stride_h2, bool swz, typename TileT>
97static __device__ __forceinline__ void load_ldmatrix(TileT & t, const half2 * tile_base, const int off_h2) {
98    if constexpr (swz) {
99        load_ldmatrix<stride_h2, swz>(t, tile_base, off_h2 / stride_h2, off_h2 % stride_h2);
100    } else {
101        ggml_cuda_mma::load_ldmatrix(t, tile_base + off_h2, stride_h2);
102    }
103}
104
105template<int stride_h2, bool swz, typename TileT>
106static __device__ __forceinline__ void load_ldmatrix_trans(
107        TileT & t, const half2 * tile_base, const int base_row, const int base_col_h2) {
108    if constexpr (swz) {
109        static_assert(std::is_same_v<TileT, ggml_cuda_mma::tile<16, 8, half2>>,
110            "the swizzled layout is only supported for tile<16, 8, half2>");
111        ldmatrix_x4_trans((int *) t.x, lane_addr<stride_h2>(tile_base, base_row, base_col_h2, TileT::I, TileT::J));
112    } else {
113        ggml_cuda_mma::load_ldmatrix_trans(t, tile_base + base_row*stride_h2 + base_col_h2, stride_h2);
114    }
115}
116
117template<int stride_h2, bool swz, typename TileT>
118static __device__ __forceinline__ void load_ldmatrix_trans(TileT & t, const half2 * tile_base, const int off_h2) {
119    if constexpr (swz) {
120        load_ldmatrix_trans<stride_h2, swz>(t, tile_base, off_h2 / stride_h2, off_h2 % stride_h2);
121    } else {
122        ggml_cuda_mma::load_ldmatrix_trans(t, tile_base + off_h2, stride_h2);
123    }
124}
125
126} // namespace ggml_cuda_fattn_smem_swizzle