DeepSeek/FlashMLA
1
0
Fork 0
mirror of https://github.com/deepseek-ai/FlashMLA synced 2025-06-26 18:15:54 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Move flash_mla.h to kernels/params.h

Browse Source
This commit is contained in:
Shengyu Liu
2025-04-22 17:46:35 +08:00
parent c7123cb36e
commit 828a19c720
7 changed files with 6 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
csrc/kernels/get_mla_metadata.h
View File

@@ -1,5 +1,5 @@
#pragma once
#include "flash_mla.h"
#include "params.h"
void run_get_mla_metadata_kernel(Mla_metadata_params &params, cudaStream_t stream);

2
csrc/kernels/mla_combine.cu
View File

@@ -5,7 +5,7 @@
#include <cutlass/array.h>
#include <cutlass/numeric_types.h>
#include "flash_mla.h"
#include "params.h"
#include "utils.h"
#include "config.h" // for BLOCK_SIZE_M and HEAD_DIM_V

2
csrc/kernels/mla_combine.h
View File

@@ -1,6 +1,6 @@
#pragma once
#include "flash_mla.h"
#include "params.h"
template<typename ElementT>
void run_flash_mla_combine_kernel(Flash_fwd_mla_params &params, cudaStream_t stream);

58
csrc/kernels/params.h Normal file
View File

@@ -0,0 +1,58 @@
#pragma once
////////////////////////////////////////////////////////////////////////////////////////////////////
struct Flash_fwd_mla_params {
using index_t = int64_t;
int b; // batch size
int s_q;
int q_seq_per_hk; // The number of q(s) per KV head, = h_q / h_k * s_q
int d, d_v; // K/V dimension
int h_q, h_k; // The number of Q/K heads
int num_blocks; // Number of blocks in total
int q_head_per_hk; // The number of q_head(s) per KV head, = h_q / h_k
bool is_causal;
float scale_softmax, scale_softmax_log2;
void *__restrict__ q_ptr;
void *__restrict__ k_ptr;
void *__restrict__ o_ptr;
void *__restrict__ softmax_lse_ptr;
index_t q_batch_stride;
index_t k_batch_stride;
index_t o_batch_stride;
index_t q_row_stride;
index_t k_row_stride;
index_t o_row_stride;
index_t q_head_stride;
index_t k_head_stride;
index_t o_head_stride;
int *__restrict__ block_table;
index_t block_table_batch_stride;
int page_block_size;
int *__restrict__ seqlens_k_ptr;
int *__restrict__ tile_scheduler_metadata_ptr;
int num_sm_parts;
int *__restrict__ num_splits_ptr;
int total_num_splits;
void *__restrict__ softmax_lseaccum_ptr;
void *__restrict__ oaccum_ptr;
};
static constexpr int TileSchedulerMetaDataSize = 8;
// [begin_idx, begin_seqlen, end_idx, end_seqlen, begin_n_split_idx, _, _, _]
struct Mla_metadata_params {
int *__restrict__ seqlens_k_ptr;
int *__restrict__ tile_scheduler_metadata_ptr;
int *__restrict__ num_splits_ptr;
int batch_size;
int block_size_n;
int fixed_overhead_num_blocks;
int num_sm_parts;
};

2
csrc/kernels/splitkv_mla.cu
View File

@@ -1,6 +1,6 @@
#include <cutlass/cutlass.h>
#include "flash_mla.h"
#include "params.h"
#include "utils.h"
#include "config.h"
#include "traits.h"

2
csrc/kernels/splitkv_mla.h
View File

@@ -1,6 +1,6 @@
#pragma once
#include "flash_mla.h"
#include "params.h"
template<typename InputT>
void run_flash_splitkv_mla_kernel(Flash_fwd_mla_params &params, cudaStream_t stream);