Add scalable sparse map kernel for large MoE models (128+ experts)

[tscholak]

Problem

The original sparse_map_kernel uses one-hot encoding to perform histogram and sorting operations:

expert_one_hot = (expert_index[:, None] == expert_range[None, :]).to(dtype)

For GPT-OSS with 128 experts and block_size=1024, this creates a [1024, 128] intermediate tensor requiring 131,072 register elements. This exceeds Triton's register pressure limits, causing the kernel to fail for num_experts > 32.

Solution

This PR introduces a new multi-pass atomic kernel that avoids large intermediate tensors:

Pass 1: Histogram with Atomics

# Each thread atomically increments its expert's counter
# Memory: O(num_experts) instead of O(block_size × num_experts)
for i in range(block_size):
    expert_id = load(top_experts_ptr + i)
    atomic_add(expert_counts_ptr + expert_id, 1)

Pass 2: Two Assignment Variants

Variant A: Atomic Assignment (faster, recommended for ≤128 experts)

• Each token atomically claims the next slot for its expert
• Best for GPT-OSS with 128 experts

Variant B: Chunked Assignment (more scalable for >128 experts)

• Process experts in chunks (e.g., 32 at a time)
• Lower memory pressure for very large expert counts

Performance

Approach	num_experts	Memory	Speed	Status
Original one-hot	≤ 32	High	Fast	✅ Works
Atomic (new)	32-128	Medium	Fastest	✅ Recommended
Chunked (new)	128+	Low	Moderate	✅ Most scalable
PyTorch fallback	Any	Low	Slow	✅ Reference

Testing

Comprehensive test suite validates correctness:

1. Small Experts (≤32): Match Original Kernel

@pytest.mark.parametrize("num_experts", [4, 8, 16, 32])
def test_scalable_kernel_matches_original_small(...)

2. Large Experts (>32): Match PyTorch Fallback

@pytest.mark.parametrize("num_experts", [64, 96, 128, 256])
def test_scalable_kernel_matches_pytorch_large(...)

3. GPT-OSS Specific Config

def test_gpt_oss_config(num_experts=128, experts_per_token=4)

4. Correctness Validation

For every test, we verify:

• ✅ All sparse_rows are unique (no collisions)
• ✅ Sparse rows within correct expert ranges
• ✅ Histogram counts match expected values
• ✅ Expert ranges are non-overlapping
• ✅ Deterministic results across runs

Run tests with:

pytest tests/functional/test_sparse_map_scalable.py -v

Files Changed

New Files

• fast_llm/functional/triton/sparse_copy_scalable.py: New kernel implementation

  • sparse_map_histogram_kernel: Atomic histogram computation
  • sparse_map_assign_kernel: Atomic index assignment
  • sparse_map_assign_chunked_kernel: Chunked index assignment
  • get_sparse_map_scalable(): Main API function

• tests/functional/test_sparse_map_scalable.py: Comprehensive test suite

  • Tests for small/large num_experts
  • GPT-OSS configuration tests
  • Edge case validation
  • Determinism tests

Usage

from fast_llm.functional.triton.sparse_copy_scalable import get_sparse_map_scalable

# For GPT-OSS with 128 experts
expert_ends, expert_pad_begins, sparse_rows = get_sparse_map_scalable(
    top_experts,
    num_experts=128,
    use_atomic_assign=True,  # Fastest for 128 experts
)

Next Steps

After this PR is merged, a follow-up PR will:

1. Integrate automatic kernel selection in get_sparse_map()
2. Update MixtureOfExpertMLP to use scalable kernel for num_experts > 32
3. Remove the dropless MoE limitation from documentation

Related

This PR builds on #374 (GPT-OSS converter) which requires 128 experts support.

🤖 Generated with Claude Code

[jlamypoirier]

Thanks for the work. Have you tried it for small number of experts? The existing kernal has disappointing performance, so not sure we stilll need it.

[jlamypoirier]

Better leave on gpu. Moving to cpu is a lot slower and needs cuda sync (really bad). You can use a @torch.compile function to avoid lots of kernles.

[jlamypoirier]

Global constant?