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

fix: handle 3D swa_indices and correct kv_bf16 expand dims

Browse Source
This commit is contained in:
biondizzle
2026-05-20 01:36:27 +00:00
parent b3f6f260ce
commit 039a9e27d6
1 changed files with 5 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
cutedsl/native_swa_decode.py
View File

@@ -94,6 +94,10 @@ def _fallback_batched_sdp(
num_tokens, NH, HD = q.shape
device = q.device
# swa_indices may be 3D (batch, tokens, window) — squeeze batch dim
if swa_indices.dim() == 3:
swa_indices = swa_indices.squeeze(0) # (tokens, window)
safe_indices = swa_indices[:num_tokens].clamp(min=0)
block_indices = safe_indices // block_size
offsets = safe_indices % block_size
@@ -116,7 +120,7 @@ def _fallback_batched_sdp(
# Batched SDPA
q_t = q.permute(1, 0, 2)
q_batch = q_t.reshape(NH * num_tokens, 1, HD)
kv_expanded = kv_bf16.unsqueeze(0).expand(NH, -1, -1, -1)
kv_expanded = kv_bf16.unsqueeze(0).expand(NH, num_tokens, window_size, HD)
k_batch = kv_expanded.reshape(NH * num_tokens, window_size, HD)
v_batch = k_batch