Restore A/B split + gsa scalar fix (error is pre-existing, not regression)

dsv4/layers/grouped_linear.py

@@ -338,10 +338,12 @@ class Nvfp4GroupedLinear:
             # gsa_gpu is (G*T,) — all rows share same amax (from max over full tensor)
             # For the GEMM's global_scale_a, fill all group slots with the same gsa value
             # Use GPU-only copy: no .item(), no CPU sync
-            self._gsa_buf[:1].copy_(gsa_gpu[:1])  # GPU→GPU scalar copy, no sync
+            self._gsa_buf[0] = gsa_gpu[0]  # scalar GPU→GPU, no sync, graph-capturable
             # Broadcast to all groups (all get same gsa)
+            # Use scalar broadcast assignment instead of copy_ from expanded view
+            # (expanded views can cause cudaErrorInvalidValue in copy_)
             if self.n_local_groups > 1:
-                self._gsa_buf[1:].copy_(self._gsa_buf[:1].expand(self.n_local_groups - 1))
+                self._gsa_buf[1:] = self._gsa_buf[0]  # scalar broadcast, graph-capturable
         else:
             self._gsa_buf.fill_(self._activation_global_scale)
             x_fp4_flat, x_sf_flat = quantize_activation_nvfp4(

dsv4/layers/linear.py

@@ -206,7 +206,7 @@ class Nvfp4Linear:
         if getattr(self, '_use_runtime_gsa', False):
             from dsv4.ops.quantize import quantize_nvfp4_gpu_fused
             x_fp4, x_sf, gsa_gpu = quantize_nvfp4_gpu_fused(hidden_states)
-            self._gsa_buf.copy_(gsa_gpu[:1].reshape(1))  # GPU → GPU, no sync
+            self._gsa_buf[0] = gsa_gpu[0]  # scalar GPU→GPU, no sync, graph-capturable
         else:
             # P2 FIX: No per-call fill_(). The _gsa_buf already has the correct
             # value — set either during initialization (via _ensure_buffer_size)
@@ -284,13 +284,10 @@ class Nvfp4Linear:
         # For M=1 decode: per-row gsa is already scalar, no reduction needed.
         # For M>1 prefill: reduce per-row gsa to a single scalar (max).
         if quant.gsa.shape[0] == 1:
-            gsa = quant.gsa[:1].reshape(1)  # Already scalar
+            self._gsa_buf[0] = quant.gsa[0]  # scalar GPU→GPU, graph-capturable
         else:
             # Reduce per-row gsa to scalar (max) for GEMM compatibility.
-            # Per-row gsa is mathematically more precise, but the GEMM only
-            # supports a single global scale per expert.
-            gsa = quant.gsa.max().reshape(1)
-        self._gsa_buf.copy_(gsa)
+            self._gsa_buf[0] = quant.gsa.max()  # GPU max, scalar assign, graph-capturable
 
         # Run GEMM
         out = run_nvfp4_grouped_gemm(

dsv4/layers/moe.py

@@ -630,7 +630,7 @@ class Nvfp4MoE:
         if getattr(self, '_use_runtime_gsa', False):
             from dsv4.ops.quantize import quantize_nvfp4_gpu_fused
             slot_x_fp4, slot_x_sf, gsa_l1_gpu = quantize_nvfp4_gpu_fused(slot_hidden)
-            self._l1_gsa_buf.copy_(gsa_l1_gpu[:1].reshape(1))  # GPU → GPU, no sync
+            self._l1_gsa_buf[0] = gsa_l1_gpu[0]  # scalar GPU→GPU, no sync, graph-capturable
         else:
             slot_x_fp4, slot_x_sf = quantize_nvfp4_gpu(
                 slot_hidden, self._l1_activation_global_scale
@@ -666,7 +666,7 @@ class Nvfp4MoE:
                 from dsv4.ops.quantize import deinterleave_amax_quantize_nvfp4_fused
                 slot_l2_x_fp4, slot_l2_x_sf, gsa_l2_gpu = deinterleave_amax_quantize_nvfp4_fused(
                     l1_out_real, self.intermediate_size)
-                self._l2_gsa_buf.copy_(gsa_l2_gpu[:1].reshape(1))  # GPU → GPU, no sync
+                self._l2_gsa_buf[0] = gsa_l2_gpu[0]  # scalar GPU→GPU, no sync, graph-capturable
             else:
                 slot_l2_x_fp4, slot_l2_x_sf = deinterleave_quantize_nvfp4_cuda(
                     l1_out_real, self.intermediate_size, self._l2_activation_global_scale
@@ -694,7 +694,7 @@ class Nvfp4MoE:
         if not self._fused_swiglu and getattr(self, '_use_runtime_gsa', False):
             from dsv4.ops.quantize import quantize_nvfp4_gpu_fused
             slot_l2_x_fp4, slot_l2_x_sf, gsa_l2_gpu = quantize_nvfp4_gpu_fused(activated)
-            self._l2_gsa_buf.copy_(gsa_l2_gpu[:1].reshape(1))  # GPU → GPU, no sync
+            self._l2_gsa_buf[0] = gsa_l2_gpu[0]  # scalar GPU→GPU, no sync, graph-capturable
         elif not self._fused_swiglu:
             slot_l2_x_fp4, slot_l2_x_sf = quantize_nvfp4_gpu(
                 activated, self._l2_activation_global_scale

dsv4/layers/shared_expert.py

@@ -268,7 +268,7 @@ class Nvfp4SharedExpert:
         if getattr(self, '_use_runtime_gsa', False):
             from dsv4.ops.quantize import quantize_nvfp4_gpu_fused
             x_fp4, x_sf, gsa_l1_gpu = quantize_nvfp4_gpu_fused(x_bf16)
-            self._l1_gsa_buf.copy_(gsa_l1_gpu[:1].reshape(1))  # GPU → GPU
+            self._l1_gsa_buf[0] = gsa_l1_gpu[0]  # scalar GPU→GPU, no sync, graph-capturable
         else:
             from dsv4.ops.quantize import quantize_activation_nvfp4
             x_fp4, x_sf = quantize_activation_nvfp4(x_bf16, self._l1_activation_global_scale)
@@ -316,7 +316,7 @@ class Nvfp4SharedExpert:
         if getattr(self, '_use_runtime_gsa', False):
             from dsv4.ops.quantize import quantize_nvfp4_gpu_fused
             x_fp4, x_sf, gsa_l1_gpu = quantize_nvfp4_gpu_fused(hidden_states)
-            self._l1_gsa_buf.copy_(gsa_l1_gpu[:1].reshape(1))  # GPU → GPU, no sync
+            self._l1_gsa_buf[0] = gsa_l1_gpu[0]  # scalar GPU→GPU, no sync, graph-capturable
         else:
             x_fp4, x_sf = quantize_activation_nvfp4(
                 hidden_states, self._l1_activation_global_scale
@@ -363,8 +363,21 @@ class Nvfp4SharedExpert:
         # Fused amax + quantize: zero CPU syncs.
         if getattr(self, '_use_runtime_gsa', False):
             from dsv4.ops.quantize import quantize_nvfp4_gpu_fused
-            x_fp4, x_sf, gsa_l2_gpu = quantize_nvfp4_gpu_fused(intermediate)
-            self._l2_gsa_buf.copy_(gsa_l2_gpu[:1].reshape(1))  # GPU → GPU, no sync
+            if not intermediate.is_contiguous():
+                intermediate = intermediate.contiguous()
+            # DEBUG: sync before quantize to isolate which kernel fails
+            torch.cuda.synchronize()
+            try:
+                x_fp4, x_sf, gsa_l2_gpu = quantize_nvfp4_gpu_fused(intermediate)
+            except RuntimeError as e:
+                print(f"  SE L2 quantize FAILED: {e}", flush=True)
+                print(f"  intermediate: shape={tuple(intermediate.shape)} dtype={intermediate.dtype} dev={intermediate.device} contiguous={intermediate.is_contiguous()}", flush=True)
+                raise
+            torch.cuda.synchronize()  # DEBUG: catch async errors from quantize
+            # Copy first element of gsa (scalar for single-expert) to pre-allocated buffer.
+            # Using scalar assignment avoids copy_() from view which caused cudaErrorInvalidValue
+            # on non-contiguous gsa_gpu slices (gsa_gpu[:1].reshape(1) — view of expanded tensor).
+            self._l2_gsa_buf[0] = gsa_l2_gpu[0]  # scalar GPU → GPU, no sync, graph-capturable
         else:
             x_fp4, x_sf = quantize_activation_nvfp4(
                 intermediate, self._l2_activation_global_scale