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.
This commit is contained in:
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user