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

D5a: Fix - add normalize param to __init__

Browse Source
This commit is contained in:
biondizzle
2026-05-23 21:11:37 +00:00
parent c80f223d08
commit d6ea7f3ebd
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -15,13 +15,14 @@ import math
class FmhaKernel:
def __init__(self, head_dim=64, s_k=128, scale_softmax=None, use_smem_p=None):
def __init__(self, head_dim=64, s_k=128, scale_softmax=None, use_smem_p=None, normalize=True):
self.head_dim = head_dim
self.s_k = s_k
self.n_kv_tiles = s_k // 128
self.pv_n_tile = min(head_dim, 256) # tcgen05 MMA max N=256
self.n_pv_tiles = head_dim // self.pv_n_tile
self.use_smem_p = use_smem_p if use_smem_p is not None else (head_dim > 64)
self.normalize = normalize # D5a: False = emit un-normalized O + lse
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