biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: update launch wrapper for fmha_decode_ref

Browse Source
This commit is contained in:
biondizzle
2026-05-28 05:28:49 +00:00
parent 00ac46c9d3
commit 77fa34a9a6
1 changed files with 5 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
dsv4/kernels/attention/fmha_sm100_launch.cu
View File

@@ -36,7 +36,7 @@ std::tuple<torch::Tensor, torch::Tensor> fmha_decode_cuda(
const float* sp = attn_sink.has_value() ? attn_sink->data_ptr<float>() : nullptr;
#define L(D, C, S) fmha_decode<D,C,S><<<grid,block,smem>>>( \
#define L(D, C, S) fmha_decode_ref<D><<<grid,block,smem>>>( \
(uint16_t*)q.data_ptr<at::BFloat16>(), \
(uint16_t*)k.data_ptr<at::BFloat16>(), \
(uint16_t*)v.data_ptr<at::BFloat16>(), \
@@ -44,10 +44,10 @@ std::tuple<torch::Tensor, torch::Tensor> fmha_decode_cuda(
q.stride(0), k.stride(0), o.stride(0), \
sk, n_comp, swa_len, (float)scale, sp, lse.data_ptr<float>())
if (D==64) { if(is_causal&&sp) L(64,1,1); else if(is_causal) L(64,1,0); else if(sp) L(64,0,1); else L(64,0,0); }
else if (D==128) { if(is_causal&&sp) L(128,1,1); else if(is_causal) L(128,1,0); else if(sp) L(128,0,1); else L(128,0,0); }
else if (D==256) { if(is_causal&&sp) L(256,1,1); else if(is_causal) L(256,1,0); else if(sp) L(256,0,1); else L(256,0,0); }
else if (D==512) { if(is_causal&&sp) L(512,1,1); else if(is_causal) L(512,1,0); else if(sp) L(512,0,1); else L(512,0,0); }
if (D==64) L(64, 0, 0);
else if (D==128) L(128, 0, 0);
else if (D==256) L(256, 0, 0);
else if (D==512) L(512, 0, 0);
else { TORCH_CHECK(false, "Unsupported head_dim: ", D); }
#undef L
return {o, lse};