From d8e17d70c14e90d691363f6c5cfce4abc6f8044d Mon Sep 17 00:00:00 2001 From: biondizzle Date: Tue, 2 Jun 2026 07:57:39 +0000 Subject: [PATCH] P0+P1+P2: Enable fused SwiGLU (MoE+SE), fix SE _run_l1_fused, remove per-call gsa fill_ MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit P0: Enable fused SwiGLU for MoE (set_fused_swiglu(True)) - Saves 240+ unfused BF16 kernel launches per token - SiLU + clamp in kernel registers instead of separate launches P1: Fix shared expert _run_l1_fused + enable fused SwiGLU - Fixed: _l1_sf_view -> _l1_scale_b, _l1_gs_view -> _l1_gsb - Fixed: expert_offsets dtype int64 -> int32 - Added proper padded buffer + scale assembly (matching unfused path) - Added runtime gsa support (quantize_nvfp4_gpu_fused) P2: Remove per-call gsa_buf.fill_() in Nvfp4Linear - fill_() was H2D transfer every forward pass (~5µs × 244 calls = ~1.2ms/token) - _gsa_buf now initialized with _activation_global_scale (not zeros) - After warmup_gsa, buffer already has correct value — no fill needed --- dsv4/layers/linear.py | 9 ++++++-- dsv4/layers/shared_expert.py | 40 ++++++++++++++++++++++++++++-------- single_shot_inference.py | 9 ++++---- 3 files changed, 43 insertions(+), 15 deletions(-) diff --git a/dsv4/layers/linear.py b/dsv4/layers/linear.py index fb007010..a7a1554b 100644 --- a/dsv4/layers/linear.py +++ b/dsv4/layers/linear.py @@ -113,7 +113,7 @@ class Nvfp4Linear: ).view(torch.float4_e2m1fn_x2) self._expert_offsets_buf = torch.zeros(1, dtype=torch.int32, device=self.device) - self._gsa_buf = torch.zeros(1, dtype=torch.float32, device=self.device) + self._gsa_buf = torch.full((1,), self._activation_global_scale, dtype=torch.float32, device=self.device) def _ensure_initialized(self): if self._mat_b is None: @@ -176,8 +176,13 @@ class Nvfp4Linear: 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 else: + # P2 FIX: No per-call fill_(). The _gsa_buf already has the correct + # value — set either during initialization (via _ensure_buffer_size) + # or by the first GPU compute when _use_runtime_gsa was True. + # Old path: self._gsa_buf.fill_(self._activation_global_scale) + # — H2D transfer every call (~5µs each × 244 calls = ~1.2ms/token). + # New path: zero H2D transfers on the hot path. from dsv4.ops.quantize import quantize_nvfp4_gpu - self._gsa_buf.fill_(self._activation_global_scale) x_fp4, x_sf = quantize_nvfp4_gpu(hidden_states, self._activation_global_scale) # Scatter x_fp4 into padded buffer diff --git a/dsv4/layers/shared_expert.py b/dsv4/layers/shared_expert.py index c770cf76..a1800ecc 100644 --- a/dsv4/layers/shared_expert.py +++ b/dsv4/layers/shared_expert.py @@ -248,21 +248,43 @@ class Nvfp4SharedExpert: num_tokens = hidden_states.shape[0] x_bf16 = hidden_states.reshape(num_tokens, self.hidden_size) - # Quantize activation to NVFP4 - x_fp4, x_sf, gsa = quantize_nvfp4_gpu_fused(x_bf16) + # Quantize activation to NVFP4 (fused amax + quantize) + 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 + else: + from dsv4.ops.quantize import quantize_activation_nvfp4 + x_fp4, x_sf = quantize_activation_nvfp4(x_bf16, self._l1_activation_global_scale) - # Run fused grouped GEMM with 1 group + # Padded buffer setup for 1-group GEMM + padded_rows = cutedsl_ceil_div(num_tokens, 128) * 128 + padded_x_fp4 = self._padded_x_fp4_buf_l1 + padded_x_fp4.view(torch.uint8).zero_() + padded_x_fp4.view(torch.uint8)[:num_tokens] = x_fp4.view(torch.uint8) + + # Assemble A-side scales + scale_a = self._assemble_scales_single_group(x_sf, num_tokens, self._padded_x_sf_buf_l1) + + # Expert offsets: [padded_rows] for 1 group (int32, pre-allocated) + expert_offsets = self._expert_offsets_buf + expert_offsets.fill_(padded_rows) + + # Global scales — GPU-computed gsa already in _l1_gsa_buf (no CPU sync) + gsa = self._l1_gsa_buf + + # Run fused GEMM + SwiGLU l1_out = run_fused_swiglu_grouped_gemm( - mat_a=x_fp4, + mat_a=padded_x_fp4, mat_b=self._l1_mat_b, - scale_a=x_sf, - scale_b=self._l1_sf_view, - expert_offsets=torch.tensor([num_tokens], dtype=torch.int64, device=x_fp4.device), + scale_a=scale_a, + scale_b=self._l1_scale_b, + expert_offsets=expert_offsets, global_scale_a=gsa, - global_scale_b=self._l1_gs_view, + global_scale_b=self._l1_gsb, swiglu_limit=self.swiglu_limit if self.swiglu_limit is not None else 0.0, ) - return l1_out # (num_tokens, intermediate_size) BF16, SwiGLU already applied + return l1_out[:num_tokens] # (num_tokens, intermediate_size) BF16, SwiGLU already applied def _run_l1(self, hidden_states: torch.Tensor) -> torch.Tensor: """L1 GEMM: activation × gate_up_weight → BF16.""" diff --git a/single_shot_inference.py b/single_shot_inference.py index 8b3b4ffe..d0c936e5 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -1021,7 +1021,9 @@ def main(): intermediate_size=cfg.get("moe_intermediate_size", 3072), top_k=cfg.get("num_experts_per_tok", 6), device=dev) moe.set_swiglu_limit(cfg.get("swiglu_limit", 10.0)) - moe._fused_swiglu = False # P0: Fused SwiGLU kernel has CuTeDSL arg-binding issue — disabled until kernel fix + # P0: ENABLE fused SwiGLU — NVFP4 GEMM + SiLU in kernel registers. + # Saves 240+ unfused BF16 kernel launches per token (gate_silu, clamp, mul, quantize). + moe.set_fused_swiglu(True) _load_moe_weights_stacked(all_w, li, pfx, dev, moe, cfg) # EAGERLY process stacked weights → K-major + swizzle, free raw tensors moe._ensure_stacked() @@ -1036,9 +1038,8 @@ def main(): se = Nvfp4SharedExpert(hidden_size=H, intermediate_size=cfg.get("moe_intermediate_size", 3072), device=dev, swiglu_limit=cfg.get("swiglu_limit", 10.0)) _load_shared_expert_weights(all_w, li, pfx, dev, se, cfg) - # P1: Enable fused SwiGLU for shared expert (1-group variant of MoE fused kernel) - # DISABLED: Same CuTeDSL arg-binding issue as MoE fused kernel - se.set_fused_swiglu(False) + # P1: ENABLE fused SwiGLU for shared expert (1-group variant of MoE fused kernel) + se.set_fused_swiglu(True) # EAGERLY process shared expert weights se._ensure_initialized() # Fix activation global scales — _ensure_initialized sets gsa from l1_gs (which is 1.0)