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)