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

fix: add is_causal to FmhaKernel __init__ signature

Browse Source
This commit is contained in:
biondizzle
2026-05-26 10:53:14 +00:00
parent 841a3e87b2
commit df84420414
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
dsv4/kernels/attention/fmha.py
View File

@@ -16,7 +16,7 @@ import math
class FmhaKernel:
def __init__(self, head_dim=64, s_k=128, scale_softmax=None, use_smem_p=None, normalize=True, num_query_heads=1, batch_size=1, apply_swa_mask=False):
def __init__(self, head_dim=64, s_k=128, scale_softmax=None, use_smem_p=None, normalize=True, num_query_heads=1, batch_size=1, apply_swa_mask=False, is_causal=False):
self.head_dim = head_dim
self.s_k = s_k
self.n_kv_tiles = s_k // 128
@@ -32,7 +32,7 @@ class FmhaKernel:
self.batch_size = batch_size
self.normalize = normalize # D5a: False = emit un-normalized O + lse
self.apply_swa_mask = apply_swa_mask # D3: mask logits at positions >= swa_lens
self.is_causal = False # D4: causal mask (k_coord > m_coord) on SWA branch
self.is_causal = is_causal # D4: causal mask (k_coord > m_coord) on SWA branch
self.acc_dtype = Float32; self.qk_acc_dtype = Float32
self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16
self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1