diff --git a/dsv4/layers/linear.py b/dsv4/layers/linear.py
index 4bc24ec9..52e7a626 100644
--- a/dsv4/layers/linear.py
+++ b/dsv4/layers/linear.py
@@ -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
diff --git a/dsv4/layers/shared_expert.py b/dsv4/layers/shared_expert.py
index f00aabb4..afcef89a 100644
--- a/dsv4/layers/shared_expert.py
+++ b/dsv4/layers/shared_expert.py
@@ -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)