# grouped-moe-gemv

FlashRT native CUDA grouped expert projection kernels for Blackwell decode and
small verify batches. Version 2 adds W4A4 with device-side top-k routing while
preserving the version 1 W4A16 APIs.

## Hardware Backends

- SM110 (Jetson AGX Thor): W4A16 decode and grouped expert GEMV use the
  FlashRT edge backend validated by FlashRT PR #169. This target is compiled
  independently with `FLASHRT_W4A16_EDGE_UNROLL=2`; the SM120 value remains 4.
- SM120/SM121: W4A16 and block-scaled-MMA W4A4 paths are available.
- W4A4 is intentionally rejected on SM110 because that implementation uses
  the SM120 block-scaled MMA path. It never silently falls back or launches an
  incompatible cubin.

## Functions

- `w4a16_decode_gemv_bf16(x_bf16, weight_packed, sfb, alpha=1.0, out=None)`
- `grouped_w4a16_gemv_bf16(activations, weight_stack, sfb_stack, alpha_stack, expert_idx, n, w_stride=None, sfb_stride=None, out=None)`
- `quantize_activations_nvfp4_bf16(activations, packed=None, sfa=None)`
- `quantize_weights_nvfp4_bf16(weights, packed=None, sfb=None)`
- `grouped_w4a4_gemv_bf16(activations_packed, weight_stack, sfa, sfb_stack, alpha_stack, expert_idx, out=None)`
- `grouped_w4a4_gemv_from_bf16(activations, weight_stack, sfb_stack, alpha_stack, expert_idx, packed=None, sfa=None, out=None)`

The grouped API runs one BF16-activation x NVFP4-weight GEMV per routed slot.
It is intended for static routed expert batches where the caller already owns
packed weights and swizzled scale-factor buffers.

On SM120/SM121, the W4A4 API accepts packed activations `[M,K/2]`, expert weights
`[E,N,K/2]`, and a contiguous device routing tensor `[M,top_k]`. It emits
`[M,top_k,N]` in one grouped compute launch. For down projections with a
different activation per routed pair, flatten to `M=routed_pairs, top_k=1`.

`K` must be divisible by 16 and `N` by 8. Target `K%64==0` shapes use tuned
SM120 paths; the remaining `K%16` shapes use a fixed-order SIMT contract path.
No atomics, host synchronization, or dynamic workspace are used by the native
ops. Pass `packed`, `sfa`, and `out` buffers to the composed helper for
allocation-free CUDA Graph capture.

## Example

```python
from kernels import get_kernel
import torch

try:
    moe = get_kernel(
        "flashrt/grouped-moe-gemv", version=2, trust_remote_code=True
    )
except TypeError:  # kernels==0.12.x compatibility
    moe = get_kernel("flashrt/grouped-moe-gemv", version=2)

M, TOP_K, E, N, K = 7, 8, 8, 1024, 2048
x = torch.randn(M, K, device="cuda", dtype=torch.bfloat16)
expert_idx = torch.randint(E, (M, TOP_K), device="cuda", dtype=torch.int32)

def sf_bytes(rows, dim):
    return ((rows + 127) // 128) * (((dim // 16) + 3) // 4) * 512

# Do this once while loading the checkpoint, not in the inference hot path.
weights_bf16 = torch.randn(E, N, K, device="cuda", dtype=torch.bfloat16)
weights_packed = torch.empty(E, N, K // 2, device="cuda", dtype=torch.uint8)
weight_sfs = torch.empty(E, sf_bytes(N, K), device="cuda", dtype=torch.uint8)
for expert in range(E):
    moe.quantize_weights_nvfp4_bf16(
        weights_bf16[expert],
        packed=weights_packed[expert],
        sfb=weight_sfs[expert],
    )
weight_alpha = torch.ones(E, device="cuda", dtype=torch.float32)

packed = torch.empty(M, K // 2, device="cuda", dtype=torch.uint8)
sfa = torch.empty(sf_bytes(M, K), device="cuda", dtype=torch.uint8)
out = torch.empty(M, TOP_K, N, device="cuda", dtype=torch.bfloat16)
y = moe.grouped_w4a4_gemv_from_bf16(
    x, weights_packed, weight_sfs, weight_alpha, expert_idx,
    packed=packed, sfa=sfa, out=out,
)
```

The example buffers may be larger than the minimum; wrappers validate storage.
For production code, derive SF sizes from the checkpoint packer metadata.

## Dispatch guidance

Use the built-artifact benchmark to dispatch rather than selecting only by
dtype. On the tested cu128 artifact W4A16 wins gate-up, W4A4 wins down verify,
and down decode is effectively tied kernel-only. A fused/upstream FP4 producer
removes the standalone quantization charge, but callers still should not assume
lower precision is automatically faster.

## Validation

```bash
python grouped-moe-gemv/tests/test_grouped_moe_gemv.py --backend source --mode full
python grouped-moe-gemv/benchmarks/benchmark.py --backend source
```