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

grouped_linear: Pre-allocate output buffer for grouped GEMM (CUDA graph capture)

Browse Source 
Add _output_buf_padded for the flat GEMM output, pass as out= parameter
to run_nvfp4_grouped_gemm to avoid per-step torch.zeros() allocation.
This commit is contained in:
biondizzle
2026-06-03 22:02:01 +00:00
parent 676fad064f
commit b32713c302
1 changed files with 13 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
dsv4/layers/grouped_linear.py
View File

@@ -223,6 +223,12 @@ class Nvfp4GroupedLinear:
self.max_num_tokens, self.n_local_groups, self.o_lora_rank,
dtype=torch.bfloat16, device=self.device
)
# Pre-allocate FLAT output buffer for grouped GEMM (graph capture)
# The GEMM produces (tokens_sum, n_dim) where n_dim = n_local_groups * o_lora_rank
self._output_buf_padded = torch.zeros(
self.max_num_tokens, self.n_local_groups * self.o_lora_rank,
dtype=torch.bfloat16, device=self.device
)
# Pre-allocate scale_a swizzle buffer for graph capture
K_sf = cutedsl_ceil_div(self.group_in_features, 16)
max_padded_rows = cutedsl_ceil_div(self.max_num_tokens, 128) * 128
@@ -377,8 +383,8 @@ class Nvfp4GroupedLinear:
# Global scales — GPU-computed gsa already in _gsa_buf (no CPU sync)
gsa = self._gsa_buf
# Run grouped GEMM
out = run_nvfp4_grouped_gemm(
# Run grouped GEMM — pass pre-allocated output buffer for CUDA graph capture
z_gem = run_nvfp4_grouped_gemm(
mat_a=padded_x_fp4,
mat_b=self._mat_b,
scale_a=scale_a,
@@ -386,22 +392,24 @@ class Nvfp4GroupedLinear:
expert_offsets=expert_offsets,
global_scale_a=gsa,
global_scale_b=self._gsb,
out=self._output_buf_padded if hasattr(self, '_output_buf_padded') else None,
)
# Extract real outputs and reshape
# GEMM output has the same layout as mat_a: groups-first with padding
# For CUDA graph capture (T=1 decode): use vectorized GPU gather — no Python loop.
# For T>1 prefill: Python loop is OK (not graph-captured).
z = self._output_buf[:num_tokens] if hasattr(self, '_output_buf') and self._output_buf is not None else torch.empty(num_tokens, self.n_local_groups, self.o_lora_rank, dtype=torch.bfloat16, device=o.device)
z_gem = z_gem if z_gem is not None else self._output_buf_padded
z = self._output_buf[:num_tokens]
if num_tokens == 1:
# Vectorized: gather_indices = [0, padded_T, 2*padded_T, ...] — GPU-only
gather_indices = self._expert_offsets_range_buf[:self.n_local_groups] * padded_rows_per_group - padded_rows_per_group
z_flat = out[gather_indices] # (n_groups, o_rank) — GPU gather
z_flat = z_gem[gather_indices] # (n_groups, o_rank) — GPU gather
z[:, :, :] = z_flat.unsqueeze(0) # (1, n_groups, o_rank)
else:
for g in range(self.n_local_groups):
offset = g * padded_rows_per_group
z[:, g, :] = out[offset:offset + num_tokens, :]
z[:, g, :] = z_gem[offset:offset + num_tokens, :]
return z