From 5bb78564f55f81999914fdf620fdbb5444347e1c Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 17 May 2026 14:01:32 +0000 Subject: [PATCH] Remove dynamic tensor allocation in scale assembly (cudagraph fix) Removed torch.zeros() call that created padded_expert_offsets during scale assembly. Now uses fixed layout computed from Python constants. Also removed dead reference to padded_expert_offsets variable. --- vllm/nvfp4_cutedsl.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/vllm/nvfp4_cutedsl.py b/vllm/nvfp4_cutedsl.py index 63235048..3fb19285 100644 --- a/vllm/nvfp4_cutedsl.py +++ b/vllm/nvfp4_cutedsl.py @@ -238,12 +238,6 @@ class CuTeDSLMoERunner: padded_x_sf = padded_x_sf_buf padded_x_sf.zero_() - # Compute padded expert offsets (each expert padded to 128 rows) - tokens_per_expert = expert_offsets[1:] - expert_offsets[:-1] - padded_rows_per_expert = ((tokens_per_expert + 127) // 128) * 128 - padded_expert_offsets = torch.zeros(num_experts + 1, dtype=torch.int32, device=x_sf.device) - padded_expert_offsets[1:] = padded_rows_per_expert.cumsum(0) - # Phase 1: Scatter x_sf into fixed-layout per-expert sections # Each expert gets max_chunks * 128 rows at offset e * max_chunks * 128. # This matches Phase 2's fixed reading pattern. @@ -276,7 +270,8 @@ class CuTeDSLMoERunner: all_flat = torch.cat([p.view(torch.uint8) for p in swizzled_parts], dim=0) all_flat = all_flat.view(torch.float8_e4m3fn) - total_padded = padded_expert_offsets[num_experts] + # Total rows = num_experts * max_chunks * 128 (fixed) + total_padded = num_experts * max_chunks * 128 return all_flat.reshape(total_padded, -1) def compute_activation_global_scales(self, hidden_states_sample, topk_weights, topk_ids):