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

switch to SDPA with sinks (better residual control)

Browse Source
This commit is contained in:
biondizzle
2026-05-31 04:38:41 +00:00
parent 3ed8f3cc44
commit c69dc51b3b
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
single_shot_inference.py
View File

@@ -369,7 +369,7 @@ def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin,
# -- FMHA: (n_h, T, hd) × (1, seq_len, hd) → (n_h, T, hd) --
q_input = q_heads.permute(1, 0, 2) # (n_h, T, hd)
# Use PyTorch SDPA for correctness verification
USE_SDPA = False # Use production FMHA kernel
USE_SDPA = True # Use SDPA with sinks for correctness
if USE_SDPA:
# Expand K/V for GQA: (1, seq_len, hd) → (n_h, seq_len, hd)
k_expanded = k_full.expand(n_h, -1, -1).contiguous() # (n_h, seq_len, hd)