Dynamic layer

[Harikeshav-R]

Add the ability to dynamically resize layer processing between inference calls for load balancing purposes

[Copilot]

Pull request overview

Adds runtime “dynamic layer” resizing for distributed inference by introducing control messages over the existing ring transports, plus several SIMD/OpenMP performance optimizations in sampling and transformer math.

Changes:

• Add control-message protocol (RESIZE_LAYERS, RESIZE_CHAIN, ACK) to reconfigure layer ranges at runtime.
• Integrate resize commands into chat (/layers, /resize, /help) and propagate resizes/ACKs through workers.
• Optimize hot-path math (temperature scaling SIMD, softmax OpenMP, attention scaling reuse, residual/SILU vectorization).

Reviewed changes

Copilot reviewed 14 out of 14 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
src/inference/main.cpp	Adds chat slash commands and a master-side ACK wait helper; sets transport packet size for control padding.
src/inference/Transport.h	Defines control message structs/types and default send_control/packet sizing support.
src/inference/TcpTransport.h / .cpp	Implements control send + nonblocking control receive (poll + MSG_PEEK) for TCP.
src/inference/UdpTransport.h / .cpp	Adds control send for UDP (inline over existing chunk protocol); nonblocking control recv stubbed.
src/inference/KernelTransport.h / .cpp	Adds control send for kernel transport; nonblocking control recv stubbed (inline handling in worker loop).
src/inference/Transformer.h	Adds runtime layer-config update helper and makes clear_kv_cache() required.
src/inference/FloatTransformer.h / .cpp	Implements KV-cache clearing; adds control handling in worker loop; performance tweaks (softmax/attention/residual/SILU).
src/inference/QuantizedTransformer.h / .cpp	Same as FloatTransformer (KV clear + control handling + perf tweaks).
src/inference/Sampler.cpp	SIMD vectorization for temperature scaling.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

wait_for_resize_ack checks elapsed time, but it calls transport->recv_prev(...) in blocking mode; if no packet arrives, the function will block indefinitely and the timeout will never trigger. Consider adding a transport-level timed/non-blocking receive (e.g., poll/select + recv, or a recv_prev_with_timeout) and using that here. Also, any non-ACK packet received here is currently discarded, which would corrupt the data stream if /resize is issued while inference traffic is in flight.

Suggested change

	// Returns true if ACK received within timeout, false otherwise.
	bool wait_for_resize_ack(Transport* transport, int timeout_ms = 5000) {
	size_t packet_size = transport->get_packet_size();
	if (packet_size == 0) {
	// Fallback: no packet size set, use control header size
	packet_size = sizeof(ControlPacketHeader);
	}
	std::vector<char> buffer(packet_size);
	auto start = std::chrono::steady_clock::now();
	while (true) {
	// Check for timeout
	auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(
	std::chrono::steady_clock::now() - start).count();
	if (elapsed >= timeout_ms) {
	std::cerr << "Warning: Resize ACK timeout after " << timeout_ms << "ms" << std::endl;
	return false;
	}
	// Receive a packet (blocking)
	transport->recv_prev(buffer.data(), packet_size);
	// Check if it's an ACK
	uint16_t magic = 0;
	std::memcpy(&magic, buffer.data(), sizeof(magic));
	if (magic == CONTROL_MAGIC) {
	ControlPacketHeader pkt{};
	std::memcpy(&pkt, buffer.data(), sizeof(pkt));
	if (pkt.msg.type == ControlMessageType::ACK) {
	return true; // Successfully received ACK
	}
	}
	// Not an ACK - this shouldn't happen in normal operation
	// but continue waiting in case of spurious packets
	}
	// Returns true if an ACK packet is received, false otherwise.
	// Note: This function performs a single blocking receive; any timeout
	// semantics must be implemented at the transport or caller level.
	bool wait_for_resize_ack(Transport* transport, int /*timeout_ms*/ = 5000) {
	size_t packet_size = transport->get_packet_size();
	if (packet_size == 0) {
	// Fallback: no packet size set, use control header size
	packet_size = sizeof(ControlPacketHeader);
	}
	std::vector<char> buffer(packet_size);
	// Receive a single packet (blocking).
	transport->recv_prev(buffer.data(), packet_size);
	// Check if it's an ACK control packet.
	uint16_t magic = 0;
	std::memcpy(&magic, buffer.data(), sizeof(magic));
	if (magic == CONTROL_MAGIC) {
	ControlPacketHeader pkt{};
	std::memcpy(&pkt, buffer.data(), sizeof(pkt));
	if (pkt.msg.type == ControlMessageType::ACK) {
	return true; // Successfully received ACK
	}
	}
	// Not an ACK or not a control packet.
	return false;

[Copilot]

In the 2-boundary /resize case, msg.end_layer can be set to a value < n_layers, which will make the worker non-tail and skip running the remaining layers. The master then computes logits immediately after receiving the activation, producing incorrect results. Consider requiring the last boundary to equal n_layers for /resize, or implementing a resize command that also reconfigures downstream workers when end_layer < n_layers.

[Copilot]

For multi-worker chain resizing, there’s no validation that the final boundary assigns a tail worker (i.e., that the last boundary equals n_layers). If no worker is marked tail, the last layers won’t be executed before the master computes logits. Consider validating boundaries.back() == n_layers (or otherwise ensuring full layer coverage and exactly one tail) before sending RESIZE_CHAIN.

[Copilot]

#include <thread> is added but not used in this file. Consider removing it to avoid unused-include warnings and keep compile units minimal.

Suggested change

#include <thread>

[Copilot]

#include <poll.h> is included here but not used. Consider removing it (or adding the intended polling-based logic) to avoid unused-include warnings.

Suggested change

#include <poll.h>