chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[sgl-kernel] support custom fp8 flashmla kernel (#13087)

Browse Source
This commit is contained in:
Fan Yin
2025-11-14 04:45:21 +08:00
committed by GitHub
parent 8779100768
commit 2966367a31
5 changed files with 214 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
sgl-kernel/include/sgl_kernel_ops.h
View File

@@ -935,3 +935,21 @@ void FMHACutlassSM100FwdRun(
std::vector<at::Tensor>
sparse_prefill_fwd(const at::Tensor& q, const at::Tensor& kv, const at::Tensor& indices, double sm_scale, int64_t d_v);
std::vector<at::Tensor> fwd_kvcache_mla_fp8(
at::Tensor& q, // batch_size x seqlen_q x num_heads x head_size
const at::Tensor& kcache, // num_blocks x page_block_size x num_heads_k x head_size (when is_fp8 is False) or
// num_blocks x num_heads_k x (page_block_size*656) (when is_fp8 is True)
const int64_t head_size_v,
const at::Tensor& seqlens_k, // batch_size
const at::Tensor& block_table, // batch_size x max_num_blocks_per_seq
const double softmax_scale,
bool is_causal,
const at::Tensor& tile_scheduler_metadata, // num_sm_parts x TileSchedulerMetaDataSize
const at::Tensor& num_splits, // batch_size + 1
const std::optional<at::Tensor>& descale_q, // None or batch_size
const std::optional<at::Tensor>& descale_k // None or batch_size
);
std::vector<at::Tensor> get_mla_decoding_metadata_dense_fp8(
at::Tensor& seqlens_k, const int64_t num_heads_per_head_k, const int64_t num_heads_k);