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

Add safety check for swizzled buffers: fall through to Python path if None

Browse Source
This commit is contained in:
biondizzle
2026-06-04 04:32:00 +00:00
parent fae61d3ef7
commit 119e6d471e
2 changed files with 9 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
dsv4/layers/linear.py
View File

@@ -170,7 +170,7 @@ class Nvfp4Linear:
view = buf[:padded_rows, :padded_cols]
# During graph capture, use CUDA swizzle kernel (Python view ops not capturable)
if torch.cuda.is_current_stream_capturing():
if torch.cuda.is_current_stream_capturing() and self._padded_x_sf_swizzled_buf is not None:
from dsv4.kernels.cuda.loader import get_cuda_module
mod = get_cuda_module("blackwell_swizzle", ["blackwell_swizzle.cu"])
swizzled_buf = self._padded_x_sf_swizzled_buf

14
dsv4/layers/shared_expert.py
View File

@@ -247,12 +247,14 @@ class Nvfp4SharedExpert:
if torch.cuda.is_current_stream_capturing():
from dsv4.kernels.cuda.loader import get_cuda_module
swizzled_buf = self._padded_x_sf_swizzled_buf_l1 if padded_x_sf_buf is self._padded_x_sf_buf_l1 else self._padded_x_sf_swizzled_buf_l2
mod = get_cuda_module("blackwell_swizzle", ["blackwell_swizzle.cu"])
mod.blackwell_swizzle_32_4_4(
view.view(torch.uint8), swizzled_buf[:padded_rows, :padded_cols].view(torch.uint8),
padded_rows, padded_cols
)
return swizzled_buf[:padded_rows, :padded_cols].reshape(padded_rows, padded_cols)
if swizzled_buf is not None:
mod = get_cuda_module("blackwell_swizzle", ["blackwell_swizzle.cu"])
mod.blackwell_swizzle_32_4_4(
view.view(torch.uint8), swizzled_buf[:padded_rows, :padded_cols].view(torch.uint8),
padded_rows, padded_cols
)
return swizzled_buf[:padded_rows, :padded_cols].reshape(padded_rows, padded_cols)
# Fall through to Python path if buffer not yet allocated
# Eager path: Python view operations
swizzled_flat = pad_and_swizzle_single(view)