diff --git a/dsv4/layers/moe.py b/dsv4/layers/moe.py index 6d7bf149..b9b327d7 100644 --- a/dsv4/layers/moe.py +++ b/dsv4/layers/moe.py @@ -164,6 +164,10 @@ class Nvfp4MoE: self._l1_gsa_buf = torch.zeros(self.num_experts, dtype=torch.float32, device=self.device) self._l2_gsa_buf = torch.zeros(self.num_experts, dtype=torch.float32, device=self.device) + # Pre-allocated tokens-per-expert buffer — replaces torch.bincount + # (bincount produces data-dependent shapes, breaks CUDA graph capture) + self._tokens_per_expert_buf = torch.zeros(self.num_experts, dtype=torch.int32, device=self.device) + # Row indices for scale assembly (max_num_tokens * top_k slots) self._row_indices_buf = torch.arange( self.max_num_tokens * self.top_k, device=self.device @@ -466,7 +470,16 @@ class Nvfp4MoE: # Quantize slot_hidden for GEMM slot_x_fp4, slot_x_sf = quantize_activation_nvfp4(slot_hidden, l1_gs) - tokens_per_expert = torch.bincount(sorted_ids, minlength=self.num_experts)[:self.num_experts].int() + # Compute tokens_per_expert — CUDA-graph-safe alternative to torch.bincount. + # torch.bincount produces data-dependent shapes (violates graph capture). + # Instead, use scatter_add_ into a pre-allocated buffer (fixed shape, GPU-only). + self._tokens_per_expert_buf.zero_() + # Pre-allocated ones buffer — avoids per-call torch.ones() allocation + n_slots = sorted_ids.shape[0] + if not hasattr(self, '_ones_buf') or self._ones_buf.shape[0] < n_slots: + self._ones_buf = torch.ones(self.max_num_tokens * self.top_k, dtype=torch.int32, device=sorted_ids.device) + self._tokens_per_expert_buf.scatter_add_(0, sorted_ids, self._ones_buf[:n_slots]) + tokens_per_expert = self._tokens_per_expert_buf[:self.num_experts] expert_offsets = self._expert_offsets_buf expert_offsets.zero_() expert_offsets[1:self.num_experts + 1] = tokens_per_expert.cumsum(0) @@ -494,7 +507,9 @@ class Nvfp4MoE: padded_expert_offsets, self._padded_x_sf_buf_l1, self._per_expert_scale_bufs_l1 ) - l1_gsa = torch.full((self.num_experts,), l1_gs, dtype=torch.float32, device=device) + # l1_gsa: pre-allocated buffer, no per-call allocation + self._l1_gsa_buf.fill_(l1_gs) + l1_gsa = self._l1_gsa_buf l1_out = run_nvfp4_grouped_gemm( mat_a=padded_x_fp4, mat_b=self._l1_mat_b,