perf(sparsecore): enable dense gather reduce for v6e via FP32 intermediate buffer

[prishajain1]

SparseCore Dense Gather Reduce for v6e

Tensor cores are primarily designed for regular matrix multiplication. Since gathering for each token by expert and subsequent multiplication with expert router weights can be an irregular operation, we leverage the dedicated SparseCore kernel for this. There already exists an implementation of this kernel which works for v7x, but not on v6e. This PR explains why it did not work on v6e, and implements the fix to enable SparseCore dense gather reduce on v6e.

---

Why did it work for v7x and not v6e?

• v7x has 16 SparseCore SIMD lanes, while v6e has only 8 SIMD lanes.
• Gemma 4 uses $top_k = 8$ experts per token.

On v7x (16 lanes):

16 lanes
----------------------------------------------------------
| 8 expert rows for token A  │ 8 expert rows for token B |
----------------------------------------------------------

Hence, we get 2 output rows for one SIMD step. These two output rows are stored as bfloat16, and packed into a 32-bit unit:

One 32-bit unit
----------------------------------------------
| 16-bit bf16 token A  │ 16-bit bf16 token B |
----------------------------------------------

• $\text{Output block size} = 2 \text{ output rows} / 2 \text{ bf16 values per unit} = 1$ (valid integer block size).

On v6e (8 lanes):

8 lanes
------------------------------
| 8 expert rows for token A  |
------------------------------

One SIMD step produces 1 output row. But when attempting to store that output directly as bfloat16, 2 values are required to fill one 32-bit packed unit:

One 32-bit unit
----------------------------------------------
| 16-bit bf16 token A  │ missing             |
----------------------------------------------

• $\text{Output block size} = 1 \text{ output row} / 2 \text{ bf16 values per unit} = 0$.
• A zero-sized output block is invalid in Pallas/Mosaic and caused a fallback to TensorCore JAX execution.

---

The Fix

We allow storing the temporary SparseCore output as FP32 (packing = 1):

One 32-bit unit
---------------------------------------------
| 32-bit value (fp32)                       |
---------------------------------------------

• Now the output block size becomes $1 / 1 = 1$, which is valid and supported on 8-lane v6e SparseCore!

Why is this lossless?

There is no difference in the order of operations; both the reference and SparseCore paths perform FP32 accumulation before casting to bfloat16:

• TensorCore path: FP32 reduction result $\to$ bfloat16 output array.
• SparseCore path: FP32 reduction result $\to$ write to FP32 kernel output buffer $\to$ FP32 output array $\to$ bfloat16 output array.

---

Benchmark Results

We tested on Gemma4 MoE model on v6e-8:

export PHASED_PROFILING_DIR="$HOME/vllm_profile"
export USE_BATCHED_RPA_KERNEL=1
export USE_BATCHED_RPA_SEQ_ON_LANE=1
export MODEL_IMPL_TYPE=flax_nnx
export NEW_MODEL_DESIGN=1

export SAMPLING_MICROBATCH_SIZE=16
export USE_DISTRIBUTED_TOPK_SAMPLING=1
export GEMMA4_REPLICATE_GLOBAL_KV_WEIGHTS=1
export BATCHED_RPA_DECODE_SLIDING_BKV_SIZE=1024


vllm serve google/gemma-4-26B-A4B-it \
  --max-model-len=4096 \
  --max-num-seqs=256 \
  --tensor-parallel-size=8 \
  --max-num-batched-tokens=4096 \
  --gpu-memory-utilization=0.85 \
  --async-scheduling \
  --kv-cache-dtype=fp8 \
  --limit-mm-per-prompt '{"image":0,"video":0,"audio":0}' \
  --block-size=128 \
  --additional-config \
  '{"sharding":{"sharding_strategy":{"tensor_parallelism":8,"expert_parallelism":1}}}'

vllm bench serve \
        --backend openai \
        --endpoint /v1/completions \
        --model google/gemma-4-26B-A4B-it \
        --port 8000 \
        --dataset-name random \
        --random-input-len 1024 \
        --random-output-len 128 \
        --request-rate inf \
        --num-prompts 500 \
        --ignore-eos

• Prefill Phase Forward Pass Latency: 99 ms → 79 ms (20% reduction)
• jit_dense_gather_reduce Kernel Execution Time: 1.198 ms → 522 µs (56% reduction)

---

Tests Added / Updated

• test_output_block_packing_gate in tests/kernels/dense_gather_reduce_test.py: Verifies that v6e_bf16_topk8 is enabled via the FP32 intermediate buffer.