From f13a81d48bc139e37051a82323acbacbe0a9875b Mon Sep 17 00:00:00 2001 From: biondizzle Date: Wed, 3 Jun 2026 17:39:20 +0000 Subject: [PATCH] CUDA graph: Fix per-call allocations in grouped_linear and quantize MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 | 23 +++++++++++++++++++---- dsv4/ops/quantize.py | 6 +++--- 2 files changed, 22 insertions(+), 7 deletions(-) diff --git a/dsv4/layers/grouped_linear.py b/dsv4/layers/grouped_linear.py index 61057e22..437395bd 100644 --- a/dsv4/layers/grouped_linear.py +++ b/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): diff --git a/dsv4/ops/quantize.py b/dsv4/ops/quantize.py index aa1f5fb7..859fcf36 100644 --- a/dsv4/ops/quantize.py +++ b/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)