vanbasten23

┌────────────────────────┬────────────────┬────────┐
│          Code          │ MMLU-Pro runs  │  Mean  │
├────────────────────────┼────────────────┼────────┤
│ base (PR 3, 11a653a10) │ 0.7860, 0.7885 │ 0.7873 │
├────────────────────────┼────────────────┼────────┤
│ PR 4 (4de631839)       │ 0.7832, 0.7876 │ 0.7854 │
└────────────────────────┴────────────────┴────────┘

vanbasten23

┌──────┬──────────────────────────────────────┬─────────────────────┐
│ Run  │                Source                │ MMLU-Pro (full set) │
├──────┼──────────────────────────────────────┼─────────────────────┤
│ main │ 8c81640f2, j-3e505983-5c47-4357-a0d9 │ 0.7834 ± 0.0036     │
├──────┼──────────────────────────────────────┼─────────────────────┤
│ PR 3 │ 11a653a10, j-1296c734-f6f0-42ae-99d8 │ 0.7860 ± 0.0036     │
└──────┴──────────────────────────────────────┴─────────────────────┘

- Difference: +0.26 points, inside the ~±0.51 noise on a difference of two runs.
- Both runs used the fused EP W4A8 kernel.

vanbasten23

┌───────┬───────────────────────────────────┬─────────────────────┐
│  Run  │              Source               │ MMLU-Pro (full set) │
├───────┼───────────────────────────────────┼─────────────────────┤
│ main  │ e9cf14895 (includes merged #3074) │ 0.7843 ± 0.0036     │
├───────┼───────────────────────────────────┼─────────────────────┤
│ #3075 │ ba8714a3c                         │ 0.7860 ± 0.0036     │
└───────┴───────────────────────────────────┴─────────────────────┘

- Difference: +0.17 points. The noise on a two-run difference is about ±0.51 points, so this is not significant.
- Both runs used the fused EP kernel: the logs show "DeepSeek-V4 MXFP4 fused EP W4A8 enabled: block-512", and both mapped the intended source.

vanbasten23

Both MMLU-Pro runs succeeded, and accuracy is unchanged within noise.

┌───────┬──────────────────────────────────┬─────────────────────┐
│  Run  │              Source              │ MMLU-Pro (full set) │
├───────┼──────────────────────────────────┼─────────────────────┤
│ main  │ 0ef9144a6                        │ 0.7852 ± 0.0036     │
├───────┼──────────────────────────────────┼─────────────────────┤
│ after │ PR 1                             │ 0.7822 ± 0.0036     │
└───────┴──────────────────────────────────┴─────────────────────┘

vanbasten23

self = <absl.logging._VerbosityFlag object at 0x73a6b4c811f0>, argument = '-x'

    def _parse(self, argument: str | _T) -> _T | None:
      """Internal parse function.
    
      It returns the parsed value, and does not modify class states.
    
      Args:
        argument: str or the correct flag value type, argument to be parsed.
    

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