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

P0 complete: Kill .item() in grouped_linear, reduce hot-path syncs

Browse Source 
- 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).
This commit is contained in:
biondizzle
2026-06-01 22:21:12 +00:00
parent 8767c263ab
commit 583ad6cfe6
2 changed files with 42 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
dsv4/layers/grouped_linear.py
View File

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

23
single_shot_inference.py
View File

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