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

single_shot: switch to head-packed FMHA dispatch (1 kernel launch vs 128)

Browse Source
This commit is contained in:
biondizzle
2026-05-31 23:33:32 +00:00
parent 80002f2efc
commit a6a8755439
1 changed files with 7 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
single_shot_inference.py
View File

@@ -277,12 +277,15 @@ class HcHead:
# Production FMHA
# =====================================================================
def _run_production_fmha(q_heads, all_kv, n_h, hd, T, seq_len, scale, dev, li, w, pfx):
from dsv4.kernels.attention.production import dsv4_attention_per_head
q = q_heads.permute(1, 0, 2).contiguous(); k = all_kv.unsqueeze(0).contiguous(); v = k.clone()
from dsv4.kernels.attention.production import dsv4_attention
# Head-packed dispatch: single kernel launch for all 128 heads (MQA: 1 KV head shared)
q = q_heads.permute(1, 0, 2).contiguous() # (n_h, T, hd)
k = all_kv.unsqueeze(0).contiguous() # (1, N, hd) — MQA single KV head
v = k.clone()
sinks = w.get(f"{pfx}.sinks"); sink_bias = None
if sinks is not None: sink_bias = sinks.to(device=dev).float().reshape(n_h)
attn_out = dsv4_attention_per_head(q=q, k=k, v=v, scale=scale, n_comp=0, sink_bias=sink_bias)
return attn_out.permute(1, 0, 2)
attn_out = dsv4_attention(q=q, k=k, v=v, scale=scale, n_comp=0, sink_bias=sink_bias)
return attn_out.permute(1, 0, 2) # (T, n_h, hd)
# =====================================================================
# Attention — ALL production kernels