biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

P0+P1+P2: Enable fused SwiGLU (MoE+SE), fix SE _run_l1_fused, remove per-call gsa fill_

Browse Source 
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
This commit is contained in:
biondizzle
2026-06-02 07:57:39 +00:00
parent 61d5e7ba53
commit d8e17d70c1
3 changed files with 43 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
dsv4/layers/linear.py
View File

@@ -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

40
dsv4/layers/shared_expert.py
View File

@@ -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."""

9
single_shot_inference.py
View File

@@ -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)