Map-Reduce and Blockwise Attention

A technical screening for deep-learning roles covered map-reduce style vector operations and context parallelization for attention, including segmented Q/K/V, blockwise attention, and robust softmax. Requirements Explain a general map-reduce pattern on simple vectors, such as deriving C = [a + d, b + e, c + f] from A = [a, b, c] and B = [d, e, f]. Discuss context parallelization with segmented Q, K, and V tensors. Explain blockwise attention at a high level. Discuss robust / numerically stable softmax in the attention framework. Notes This is closer to an ML systems oral / technical screening than a standard coding prompt. A strong answer should connect map, local reduction, global aggregation, and communication cost. For attention, the expected depth likely includes why naive softmax over blocks is numerically wrong and how to maintain running maxima and normalizers. The canonical blockwise-attention skeleton, also called online softmax: for each Q tile, stream K/V tiles through SRAM, and maintain three running statistics per row — current max m, normalizer l = Σ exp(x - m), and output accumulator o. When a new tile's max exceeds the running max, rescale the previously accumulated l and o by exp(m_old - m_new) before adding the new tile's contribution. This trades extra FLOPs for a quadratic-to-linear drop in HBM traffic, which is why it ends up faster end-to-end despite doing more arithmetic. Be ready to draw why this is exact (not approximate) attention: the rescaling identity makes the streaming computation algebraically equivalent to materializing the full softmax(QK^T)V. Preparation Review map-reduce primitives and how vector addition decomposes across partitions. Review FlashAttention-style online softmax: block maxima, exponent rescaling, and denominator accumulation. Be ready to draw tensor shapes for Q, K, V segmentation and discuss communication between context-parallel workers.