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

P8: Unified dispatch — multi-tile kernel handles all N

Browse Source 
production.py: Single fast path using multi-tile kernel for all N.
Eliminates the separate _dsv4_attention_fast_decode path.
This commit is contained in:
biondizzle
2026-05-30 17:19:09 +00:00
parent e747742598
commit 9d483b1c54
1 changed files with 3 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
dsv4/kernels/attention/production.py
View File

@@ -453,10 +453,9 @@ def dsv4_attention(
n_segments = (N + s_k_per_seg - 1) // s_k_per_seg
if T == 1 and hd in (64, 128, 256, 512):
if n_segments == 1 and hd in (64, 128, 256):
return _dsv4_attention_fast_decode(q, k, v, scale, n_comp, sink_bias)
else:
return _dsv4_attention_multitile(q, k, v, scale, n_comp, sink_bias)
# P8: Unified fast path — multi-tile kernel handles all N
# (single-tile and multi-tile, T=1 decode)
return _dsv4_attention_multitile(q, k, v, scale, n_comp, sink_bias)
# ==================================================================
# SLOW PATH: CuTeDSL kernel + Python KV merge