From 583ad6cfe62a5b58fd93916addd5a3760228b4c7 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 1 Jun 2026 22:21:12 +0000 Subject: [PATCH] P0 complete: Kill .item() in grouped_linear, reduce hot-path syncs - grouped_linear.py: Replace .item() gsa + Python quantize with quantize_nvfp4_gpu_fused (zero CPU syncs). Flatten all groups into (G*T, D), single fused kernel launch, GPU-only gsa copy. - single_shot_inference.py: Reduce torch.cuda.synchronize() to every 20 steps instead of every step. Gate per-layer diagnostics to li<3 or li>=58 (avoid 61 .item() calls per decode step). --- dsv4/layers/grouped_linear.py | 50 ++++++++++++++++++++--------------- single_shot_inference.py | 23 +++++++++------- 2 files changed, 42 insertions(+), 31 deletions(-) diff --git a/dsv4/layers/grouped_linear.py b/dsv4/layers/grouped_linear.py index c6fb1c4e..3281d73f 100644 --- a/dsv4/layers/grouped_linear.py +++ b/dsv4/layers/grouped_linear.py @@ -17,6 +17,7 @@ import torch from dsv4.ops.quantize import ( quantize_activation_nvfp4, quantize_weight_to_nvfp4, + quantize_nvfp4_gpu_fused, ) from dsv4.ops.layouts import ( make_b_k_major, @@ -293,35 +294,42 @@ class Nvfp4GroupedLinear: # Permute to groups-first: (G, T, D) o_grouped = o_grouped.permute(1, 0, 2) - # Compute activation global scale at runtime if requested. - if getattr(self, '_use_runtime_gsa', False): - amax = o.float().abs().max().clamp(min=1e-8).item() - self._activation_global_scale = amax / (6.0 * 448.0) + # Flatten all groups into (G*T, D) for batched fused quantize — single kernel launch + o_flat = o_grouped.reshape(self.n_local_groups * num_tokens, self.group_in_features) - # Quantize each group's activation and scatter into padded buffer + # Fused amax + quantize: zero CPU-GPU syncs. + # Computes gsa on GPU, quantizes to NVFP4, returns GPU tensor. + # Replaces the old path: .item() sync + Python quantize per group. + if getattr(self, '_use_runtime_gsa', False): + x_fp4_flat, x_sf_flat, gsa_gpu = quantize_nvfp4_gpu_fused(o_flat) + # 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 + # Broadcast to all groups (all get same gsa) + if self.n_local_groups > 1: + self._gsa_buf[1:].copy_(self._gsa_buf[:1].expand(self.n_local_groups - 1)) + else: + self._gsa_buf.fill_(self._activation_global_scale) + x_fp4_flat, x_sf_flat = quantize_activation_nvfp4( + o_flat, self._activation_global_scale + ) + + # Reshape FP4 back to (G, T, D//2) and scatter into padded buffer padded_x_fp4 = self._padded_x_fp4_buf padded_x_fp4.view(torch.uint8).zero_() - # We need to collect scales for ALL groups for the GEMM - all_x_sf = [] + x_fp4_grouped = x_fp4_flat.reshape(self.n_local_groups, num_tokens, self.group_in_features // 2) for g in range(self.n_local_groups): - group_act = o_grouped[g] # (T, group_in_features) - - # Quantize this group's activation - x_fp4_g, x_sf_g = quantize_activation_nvfp4( - group_act, self._activation_global_scale - ) - - # Scatter into the padded buffer at the correct offset offset = g * padded_rows_per_group - padded_x_fp4.view(torch.uint8)[offset:offset + num_tokens] = x_fp4_g.view(torch.uint8) + padded_x_fp4.view(torch.uint8)[offset:offset + num_tokens] = x_fp4_grouped[g].view(torch.uint8) - all_x_sf.append(x_sf_g) + # Reshape scales back to (G, T, D//16) and assemble + x_sf_grouped = x_sf_flat.reshape(self.n_local_groups, num_tokens, self.group_in_features // 16) + all_x_sf = [x_sf_grouped[g] for g in range(self.n_local_groups)] # Assemble A-side scales for all groups - # The grouped GEMM expects scales for all groups assembled together - # For 2Dx3D scenario, scale_a is assembled from per-group scale tensors from dsv4.ops.layouts import ( assemble_scales_2d_side, ) @@ -332,8 +340,8 @@ class Nvfp4GroupedLinear: for g in range(self.n_local_groups): expert_offsets[g] = (g + 1) * padded_rows_per_group - # Global scales (same for all groups) - gsa = self._gsa_buf.fill_(self._activation_global_scale) + # Global scales — GPU-computed gsa already in _gsa_buf (no CPU sync) + gsa = self._gsa_buf # Run grouped GEMM out = run_nvfp4_grouped_gemm( diff --git a/single_shot_inference.py b/single_shot_inference.py index 63327935..c52a9a72 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -513,7 +513,7 @@ def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin, x_in_f, ctx_f = ffn_mhc.pre_block(X_mid); x_ffn = rmsnorm(x_in_f, ffn_norm_w) F_ffn = moe_forward(x_ffn, li, moe_runner, se_runner, router, token_id) X_next = ffn_mhc.post_block(X_mid, F_ffn, ctx_f) - if VERBOSE >= 1: + if VERBOSE >= 1 and (li < 3 or li >= 58): print(f" L{li}: |X|={X_l.abs().max().item():.1f}->{X_next.abs().max().item():.1f} " f"|Fa|={F_attn.abs().max().item():.1f} |Ff|={F_ffn.abs().max().item():.1f}", flush=True) # Detailed diagnostics for last 3 layers or any layer with explosive growth @@ -980,16 +980,19 @@ def main(): x_out = hc_head.forward(X) if hc_head is not None else X[:, 0, :] if final_norm_w is not None: x_out = rmsnorm(x_out, final_norm_w) logits = lm_head_lin(x_out) - torch.cuda.synchronize() # catch CUDA errors at source + # Only sync + validate on first 3 steps and every 20th step (reduces pipeline stalls) + if step < 3 or (step + 1) % 20 == 0: + torch.cuda.synchronize() # catch CUDA errors at source ls = logits.float() - has_nan = torch.isnan(ls).any().item() - has_inf = torch.isinf(ls).any().item() - print(f" logits: shape={list(logits.shape)} dtype={logits.dtype} " - f"min={ls.min().item():.1f} max={ls.max().item():.1f} " - f"nan={has_nan} inf={has_inf}", flush=True) - if has_nan or has_inf: - print(f" NaN/Inf in logits at step {step}, aborting", flush=True) - break + if step < 3 or (step + 1) % 20 == 0: + has_nan = torch.isnan(ls).any().item() + has_inf = torch.isinf(ls).any().item() + print(f" logits: shape={list(logits.shape)} dtype={logits.dtype} " + f"min={ls.min().item():.1f} max={ls.max().item():.1f} " + f"nan={has_nan} inf={has_inf}", flush=True) + if has_nan or has_inf: + print(f" NaN/Inf in logits at step {step}, aborting", flush=True) + break # Sampling — fused CUDA kernel (or greedy argmax for temp=0) if is_greedy: next_id = torch.argmax(logits, -1).item()