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:
@@ -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
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user