vanbasten23

"""
Weight-stationary systolic array simulation of C = A @ B  (pure Python, no deps).

Model of one Processing Element (PE) at grid position (k, n):
  - holds a STATIONARY weight  W[k][n] = B[k][n]
  - each cycle:
        a_in  = activation from WEST  (external input if n==0, else left neighbor)
        p_in  = partial sum from NORTH (0 if k==0, else top neighbor)
        p_out = p_in + a_in * W[k][n]   -> registered, travels SOUTH
        a_out = a_in                    -> registered, travels EAST

vanbasten23

# Start window and pane numbering at 1
set -g base-index 1
set -g pane-base-index 1
set -g renumber-windows on

# Mouse support
set -g mouse on

# Copy selected text into the macOS clipboard
bind-key -T copy-mode MouseDragEnd1Pane \

vanbasten23

```
x = [1, 2]

a = []
a.append(x)

x.append(3)

print(a)
# [[1, 2, 3]]

vanbasten23

$ pytest -s -vv  tpu-inference/tests/layers/vllm/test_fp8.py -k test_fused_moe 
    def test_fused_moe(use_ep, num_devices, num_tokens, intermediate_size,
                       hidden_size, num_experts, topk, enable_attn_dp):
        # Skip if enable_attn_dp is True but we don't have enough devices
        if enable_attn_dp and num_devices < 2:
            pytest.skip("enable_attn_dp requires at least 2 devices")

        mesh = test_utils.get_spmd_mesh(num_devices, enable_attn_dp)

        torch.manual_seed(42)

vanbasten23

# Plan: Add tgmm_v2 and VJP to gmm_v2

## Context

We need to add backward pass (VJP) support for 'gmm_v2'. The test at 'tests/kernels/gmm_v2_test.py:136-193' already calls 'jax.vjp(gmm_v2, ...)' and expects both gradients. Currently 'gmm_v2' has no 'custom_vjp', so the backward pass fails.

All changes go in 'tpu_inference/kernels/megablox/gmm_v2.py'.

## Math: Forward and Backward Pass of GMM

vanbasten23

# AllToAll Communication Cost on ND Torus

## Setup

- **AllToAll** on an AxBxC... torus, N = A·B·C·... total nodes
- V total bytes, W\_ici = bidirectional bandwidth per link

## Key idea: per-link load determines time

In AllToAll, each node sends V/N² bytes to each other node. To find the time, we compute the **load on the most congested link** — that's the bottleneck.

vanbasten23

(EngineCore_DP0 pid=467) ERROR 01-23 03:35:58 [core.py:935] Traceback (most recent call last):
(EngineCore_DP0 pid=467) ERROR 01-23 03:35:58 [core.py:935]   File "/workspace/vllm/vllm/v1/engine/core.py", line 926, in run_engine_core
(EngineCore_DP0 pid=467) ERROR 01-23 03:35:58 [core.py:935]     engine_core = EngineCoreProc(*args, engine_index=dp_rank, **kwargs)
(EngineCore_DP0 pid=467) ERROR 01-23 03:35:58 [core.py:935]                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=467) ERROR 01-23 03:35:58 [core.py:935]   File "/workspace/vllm/vllm/v1/engine/core.py", line 691, in __init__
(EngineCore_DP0 pid=467) ERROR 01-23 03:35:58 [core.py:935]     super().__init__(
(EngineCore_DP0 pid=467) ERROR 01-23 03:35:58 [core.py:935]   File "/workspace/vllm/vllm/v1/engine/core.py", line 105, in __init__
(EngineCore_DP0 pid=467) ERROR 01-23 03:35:58 [core.py:935]     self.model_executor = executor_class(vllm_config)
(EngineCore_DP0 pid=467) ERROR 01-23 03:35:58 [core.py:935]          

vanbasten23

tests/layers/vllm/test_unquantized.py::test_fused_moe[False-silu-False-2-8-128-1024-8-1-True] FAILED
=================================== FAILURES ===================================
____________ test_fused_moe[False-silu-False-2-8-128-1024-8-1-True] ____________
use_ep = True, num_devices = 1, num_tokens = 8, intermediate_size = 1024
hidden_size = 128, num_experts = 8, topk = 2, has_bias = False
activation = 'silu', enable_attn_dp = False
    @pytest.mark.parametrize("use_ep", [True, False])
    @pytest.mark.parametrize("num_devices", [1, jax.local_device_count()])
    @pytest.mark.parametrize("num_tokens", [8])
    @pytest.mark.parametrize("intermediate_size", [1024, 2048])

vanbasten23

exec ${PAGER:-/usr/bin/less -R} "$0" || exit 1

Test settings: forge with network access
Host details: itmm4.prod.google.com Linux 6.6.65-smp-1300.170.0.0 x86_64 astoria-genoa-base
executor.INFO: analog/view?storage=borgremote&bns=/bns/it/borg/it/bns/build-forge-executor-tpu/prod-cbf-ghostlite.forge-executor/0&min_time=1764872604000000&ts=1764872614000000
Test command:
cd /build/work/aef67bf50706fee86777a93cc065340a246c/google3/runfiles/google3 && \
  env - \
    BORG_CELL=it \
    CUSTOM_METRICS_DIR=/build/work/aef67bf50706fee86777a93cc065340a246c/google3/../custom_metrics \

vanbasten23

Let's trace the values for my_id = 1 with num_devices = 4:

outer_step	phase	Accumulation Source	left_copy_device	right_copy_device	Device providing the data
0	LEFT	x_ref[left_copy_device, ...]	(1+0+1)%4 = 2	(1-0-1)%4 = 0	Device 2
0	RIGHT	x_ref[right_copy_device, ...]	(1+0+1)%4 = 2	(1-0-1)%4 = 0	Device 0
1	LEFT	x_ref[left_copy_device, ...]	(1+1+1)%4 = 3	(1-1-1)%4 = 3	Device 3
1	RIGHT	x_ref[right_copy_device, ...]	(1+1+1)%4 = 3	(1-1-1)%4 = 3	Device 3
2	LEFT	x_ref[left_copy_device, ...]	(1+2+1)%4 = 0	(1-2-1)%4 = 2	Device 0
2	RIGHT	x_ref[right_copy_device, ...]	(1+2+1)%4 = 0	(1-2-1)%4 = 2	Device 2
As you can see, with each outer_step, the *_copy_device variables change, ensuring that the reduction operation fetches data from a new, distinct device. This systematic progression guarantees that by the end of all steps, each device has accumulated its required portion of the total sum from all other devices.