CUDA graph: Fix per-call allocations in grouped_linear and quantize

1. grouped_linear.py: Pre-allocate _scale_a_buf for swizzle
   - Same fix as linear.py — avoids torch.zeros per call
   - Uses correctly-sized view for pad_and_swizzle_single

2. quantize.py: Replace torch.zeros_like with scalar 0.0
   - torch.zeros_like allocates a full tensor every call
   - torch.where(cond, 0.0, x) broadcasts scalar — no allocation

dsv4/layers/grouped_linear.py

@@ -218,11 +218,18 @@ class Nvfp4GroupedLinear:
             1, self.n_local_groups + 1, dtype=torch.int32, device=self.device
         )
         self._group_offset_buf = torch.zeros(self.n_local_groups, dtype=torch.int32, device=self.device)
-        # Pre-allocate output buffer for graph capture (T=1 decode: 1, n_groups, o_rank)
+        # Pre-allocate output buffer for graph capture
         self._output_buf = torch.zeros(
             self.max_num_tokens, self.n_local_groups, self.o_lora_rank,
             dtype=torch.bfloat16, device=self.device
         )
+        # Pre-allocate scale_a swizzle buffer for graph capture
+        K_sf = cutedsl_ceil_div(self.group_in_features, 16)
+        max_padded_rows = cutedsl_ceil_div(self.max_num_tokens, 128) * 128
+        max_padded_cols = cutedsl_ceil_div(K_sf, 4) * 4
+        self._scale_a_buf = torch.zeros(
+            max_padded_rows, max_padded_cols, dtype=torch.float16, device=self.device
+        ).to(torch.float8_e4m3fn)
         self._buffers_allocated = True
 
     def _ensure_initialized(self):
@@ -232,14 +239,22 @@ class Nvfp4GroupedLinear:
             self._allocate_buffers()
 
     def _assemble_scales_single_group(self, x_sf):
-        """Assemble 2D-side activation scales for num_groups=1."""
+        """Assemble 2D-side activation scales for num_groups=1.
+        
+        CUDA-graph-safe: uses pre-allocated _scale_a_buf.
+        """
         num_rows, num_cols = x_sf.shape
         padded_rows = cutedsl_ceil_div(num_rows, 128) * 128
         padded_cols = cutedsl_ceil_div(num_cols, 4) * 4
 
-        buf = torch.zeros(padded_rows, padded_cols, dtype=torch.float16, device=x_sf.device).to(torch.float8_e4m3fn)
+        # Use pre-allocated buffer — zero + scatter pattern (no new allocation)
+        buf = self._scale_a_buf
+        assert buf.shape[0] >= padded_rows and buf.shape[1] >= padded_cols, \
+            f"scale_a_buf too small: {buf.shape} < ({padded_rows}, {padded_cols})"
+        buf.view(torch.uint8).zero_()
         buf[:num_rows, :num_cols] = x_sf
-        swizzled_flat = pad_and_swizzle_single(buf)
+        view = buf[:padded_rows, :padded_cols]
+        swizzled_flat = pad_and_swizzle_single(view)
         return swizzled_flat.reshape(padded_rows, padded_cols)
 
     def compute_activation_global_scale(self, o_sample: torch.Tensor):

dsv4/ops/quantize.py

@@ -80,12 +80,12 @@ def quantize_to_nvfp4(x_bf16, block_size=SF_VEC_SIZE):
     zero_block = block_amax < (6.0 * 2.0 ** -9)  # < ~0.0117
     # Zero out x for zero/underflow blocks before division.
     # This ensures x_scaled = 0 → FP4 nibbles = 0.
-    x_reshaped = torch.where(zero_block.unsqueeze(-1),
-                              torch.zeros_like(x_reshaped), x_reshaped)
+    # Use scalar 0.0 instead of torch.zeros_like — no allocation, graph-safe.
+    x_reshaped = torch.where(zero_block.unsqueeze(-1), 0.0, x_reshaped)
     block_amax = block_amax.clamp(min=1e-8)
     block_scale = (block_amax / 6.0).to(torch.float8_e4m3fn)
     # Force zero/underflow blocks: FP8 scale = 0 (exact zero).
-    block_scale = torch.where(zero_block, torch.zeros_like(block_scale), block_scale)
+    block_scale = torch.where(zero_block, 0.0, block_scale)
 
     # Nearest E2M1
     block_sf_expanded = block_scale.float().unsqueeze(-1)