Fix OOB: size padded buffers for num_experts*max_chunks*128

padded_max_slots was computed from max_tokens*top_k (3072) but
total_padded_slots in run() is num_experts*max_chunks*128 (6144).
The buffer was too small, causing index out of bounds.
This commit is contained in:
2026-05-17 14:59:45 +00:00
parent 8ac8e20fa9
commit c5af1aba6b

View File

@@ -138,10 +138,9 @@ class CuTeDSLMoERunner:
self.max_num_tokens * self.top_k, device=self.device
)
# Padded hidden/activated: per-layer, sized for max capture budget
# NOT sized for max_num_tokens (8192) which would be too much for 60 layers
max_slots = self.max_num_tokens * self.top_k
padded_max_slots = ((max_slots + 127) // 128) * 128
# Padded hidden/activated: sized for num_experts * max_chunks * 128 (the actual padded size)
max_rows_per_expert = self._max_chunks_per_expert * 128
padded_max_slots = self.num_experts * max_rows_per_expert
self._padded_hidden_buf = torch.zeros(
padded_max_slots, self.hidden_size, dtype=torch.bfloat16, device=self.device
)