GEMM¶

`batch_gemm`	Batch GEMM on Blackwell SM100 tensor cores.
`dense_gemm`	Dense GEMM on Blackwell SM100 tensor cores.
`fused_swiglu_quant`	Fused SwiGLU activation + FP8/FP4 quantization in a single pass.
`group_contiguous_gemm`	Grouped contiguous GEMM for MoE on Blackwell.
`group_masked_gemm`	Grouped masked GEMM for MoE on Blackwell with per-expert token masking.
`group_wise_quant`	Group-wise quantization from BF16/FP16 to FP8/FP4 with scale factors.
`prefill_permute_to_masked_gemm`	Permute tokens into padded expert layout for masked GEMM.

batch_gemm¶

batch_gemm( input, weight, sfa, sfb, output, bias=None, alpha=None, sm_count=-1, mode='w8a8-mxfp8', )¶

Batch GEMM on Blackwell SM100 tensor cores.

Computes output[b] = input[b] @ weight[b]^T for each batch item. The batch dimension B acts as the group (expert) index: batch item b is multiplied only by weight[b], i.e. there is a one-to-one mapping between input batch index and weight batch index (no separate expert_ids tensor is needed). Supports FP8 (W8A8-MXFP8) and both FP4 (W4A4) recipes (NVFP4 and MXFP4) via CUTLASS.

Parameters:

input (torch.Tensor) –
[B, M, K] CUDA tensor.
- w8a8-mxfp8: dtype float8_e4m3fn.
- w4a4-*: dtype uint8 (packed FP4, logical K is size(2)*2).
weight (torch.Tensor) – [B, N, K] CUDA tensor. Same dtype/packing rules as input.
sfa (torch.Tensor) –
Scale factor for A, dtype uint8, CUTLASS SfAtom 128x4 swizzled layout. Shape: [ceil(B * M / 128) * 128, sf_k], where sf_k = K // 32 (w8a8-mxfp8, w4a4-mxfp4) or K // 16 (w4a4-nvfp4).
- w8a8-mxfp8, w4a4-mxfp4: UE8M0 encoding.
- w4a4-nvfp4: UE4M3 encoding.
sfb (torch.Tensor) – [B, N, sf_k] scale factor for B, dtype uint8, same SfAtom layout as sfa.
output (torch.Tensor) – [B, M, N] pre-allocated output tensor, dtype bfloat16. Written in-place.
bias (torch.Tensor | None) – Optional [B, N] bfloat16 bias vector. Only supported in w8a8-mxfp8 mode.
alpha (torch.Tensor | None) – Optional float32 per-batch scaling factor with shape [B], one value per batch item. Required for w4a4-nvfp4; not accepted for w4a4-mxfp4.
sm_count (int) – Number of SMs to use (-1 for all available).
mode (str) – "w8a8-mxfp8", "w4a4-mxfp4" or "w4a4-nvfp4".

Returns:

None. The result is written in-place to output.

Example

from halops import batch_gemm, group_wise_quant

b, m, n, k = 4, 256, 256, 256
x_bf16 = torch.randn(b, m, k, dtype=torch.bfloat16, device="cuda")
w_bf16 = torch.randn(b, n, k, dtype=torch.bfloat16, device="cuda")
x, sfa = group_wise_quant(x_bf16.view(b * m, k), mode="w8a8-mxfp8")
w, sfb = group_wise_quant(w_bf16.view(b * n, k), mode="w8a8-mxfp8")
x = x.view(b, m, k); sfa = sfa.view(b, m, -1)
w = w.view(b, n, k); sfb = sfb.view(b, n, -1)
out = torch.empty(b, m, n, dtype=torch.bfloat16, device="cuda")
batch_gemm(x, w, sfa, sfb, out, mode="w8a8-mxfp8")

dense_gemm¶

dense_gemm( input, weight, sfa, sfb, output, bias=None, alpha=None, sm_count=-1, mode='w8a8-mxfp8', use_pdl=True, )¶

Dense GEMM on Blackwell SM100 tensor cores.

Computes output = input @ weight^T + bias. Supports FP8 (W8A8-MXFP8) and FP4 (W4A4) quantized GEMM via CUTLASS, where the FP4 recipe is selected explicitly by mode.

Parameters:

input (torch.Tensor) –
[M, K] CUDA tensor.
- w8a8-mxfp8: dtype float8_e4m3fn.
- w4a4-*: dtype uint8 (packed FP4, physical shape [M, K // 2]).
weight (torch.Tensor) – [N, K] CUDA tensor. Same dtype and packing rules as input.
sfa (torch.Tensor) –
Scale factor for A, dtype uint8, CUTLASS SfAtom 128x4 swizzled layout. The layout covers whole 128-row atoms, so the shape is [ceil(M / 128) * 128, K // 32] (w8a8-mxfp8, w4a4-mxfp4) or [ceil(M / 128) * 128, K // 16] (w4a4-nvfp4) – exactly what group_wise_quant returns.
- w8a8-mxfp8, w4a4-mxfp4: UE8M0 encoding.
- w4a4-nvfp4: UE4M3 encoding.
sfb (torch.Tensor) – Scale factor for B, dtype uint8, in the same ordinary SfAtom layout as sfa. Shape [N, K // 32] (w8a8-mxfp8, w4a4-mxfp4) or [N, K // 16] (w4a4-nvfp4).
output (torch.Tensor) – [M, N] pre-allocated output tensor, dtype bfloat16. Written in-place.
bias (torch.Tensor | None) – Optional [N] bfloat16 bias vector. Only supported in w8a8-mxfp8 mode.
alpha (torch.Tensor | None) – Optional float32 scaling factor with shape [1]. Required for w4a4-nvfp4; not accepted for w4a4-mxfp4.
sm_count (int) – Number of SMs to use (-1 for all available).
mode (str) – Quantization mode: "w8a8-mxfp8", "w4a4-mxfp4" or "w4a4-nvfp4".
use_pdl (bool) – Enable Programmatic Dependent Launch (on by default).

Returns:

None. The result is written in-place to output.

Example

from halops import dense_gemm, group_wise_quant

m = n = k = 256
x_bf16 = torch.randn(m, k, dtype=torch.bfloat16, device="cuda")
w_bf16 = torch.randn(n, k, dtype=torch.bfloat16, device="cuda")
x, sfa = group_wise_quant(x_bf16, mode="w8a8-mxfp8")
w, sfb = group_wise_quant(w_bf16, mode="w8a8-mxfp8")
out = torch.empty(m, n, dtype=torch.bfloat16, device="cuda")
dense_gemm(x, w, sfa, sfb, out, mode="w8a8-mxfp8")

fused_swiglu_quant¶

fused_swiglu_quant( input, token_nums_per_expert, is_fp4=False, )¶

Fused SwiGLU activation + FP8/FP4 quantization in a single pass.

Computes SwiGLU: y = up * silu(gate) where input = [gate | up] (concatenated along dim-1), then quantizes the result to FP8 or FP4. Outputs quantized data and CUTLASS SfAtom-formatted scale factors.

Parameters:

input (torch.Tensor) – [E * M, 2 * hidden_size] bfloat16. First half of dim-1 is gate, second half is up projection.
token_nums_per_expert (torch.Tensor) – [E] int32. Valid token count per expert.
is_fp4 (bool) – True for FP4 E2M1 output (group_size=16); False for FP8 E4M3 output (group_size=32).

Returns:

A tuple (out_quant, out_scale) where –

out_quant:
- FP4: [E*M, hidden_size/2] uint8 (packed nibbles)
- FP8: [E*M, hidden_size] float8_e4m3fn
out_scale: [E*M, kb_dim] uint8, CUTLASS SfAtom 128x4 interleaved layout.
- FP8: kb_dim = ceil(hidden_size / 128 / 4) * 16
- FP4: kb_dim = hidden_size / 16

Return type:

Tuple[torch.Tensor, torch.Tensor]

Example

from halops import fused_swiglu_quant

e, m, hidden_size = 4, 128, 2048
gate_up = torch.randn(e * m, 2 * hidden_size,
                      dtype=torch.bfloat16, device="cuda")
token_nums_per_expert = torch.full((e,), m,
                                   dtype=torch.int32, device="cuda")
out_quant, out_scale = fused_swiglu_quant(
    gate_up, token_nums_per_expert, is_fp4=False)

group_contiguous_gemm¶

group_contiguous_gemm( input, weight, sfa, sfb, expert_ids, output, bias=None, alpha=None, sm_count=-1, mode='w8a8-mxfp8', )¶

Grouped contiguous GEMM for MoE on Blackwell.

Unlike masked GEMM, tokens from different experts are packed contiguously (no per-expert padding). expert_ids maps each token row to its corresponding expert weight.

Parameters:

input (torch.Tensor) –
[M, K] CUDA tensor. All tokens packed contiguously.
- w8a8-mxfp8: dtype float8_e4m3fn.
- w4a4-*: dtype uint8 (packed FP4, logical K is size(1)*2).
weight (torch.Tensor) – [num_experts, N, K] 3D CUDA tensor. Per-expert weights.
sfa (torch.Tensor) –
Scale factor for A, dtype uint8, CUTLASS SfAtom 128x4 swizzled layout. Shape: [ceil(M / 128) * 128, sf_k], where sf_k = K // 32 (w8a8-mxfp8, w4a4-mxfp4) or K // 16 (w4a4-nvfp4).
- w8a8-mxfp8, w4a4-mxfp4: UE8M0 encoding.
- w4a4-nvfp4: UE4M3 encoding.
sfb (torch.Tensor) – Scale factor for B, dtype uint8, same SfAtom layout as sfa. shape: [num_experts, N, K // 32] (w8a8-mxfp8, w4a4-mxfp4) or [num_experts, N, K // 16] (w4a4-nvfp4).
expert_ids (torch.Tensor) – [M] int32. Maps each row to an expert index in [0, num_experts). Must be sorted in non-decreasing order.
output (torch.Tensor) – [M, N] bfloat16 output. Written in-place.
bias (torch.Tensor | None) – Optional [num_experts, N] bfloat16 bias. Only supported in w8a8-mxfp8 mode.
alpha (torch.Tensor | None) – Optional float32 per-expert scale with shape [num_experts]. Required for w4a4-nvfp4; not accepted for w4a4-mxfp4.
sm_count (int) – Number of SMs to use (-1 for all available).
mode (str) – "w8a8-mxfp8", "w4a4-mxfp4" or "w4a4-nvfp4".

Returns:

None. The result is written in-place to output.

Example

from halops import group_contiguous_gemm, group_wise_quant

m, n, k = 256, 256, 256
num_experts = 4
x_bf16 = torch.randn(m, k, dtype=torch.bfloat16, device="cuda")
w_bf16 = torch.randn(num_experts, n, k, dtype=torch.bfloat16, device="cuda")
x, sfa = group_wise_quant(x_bf16, mode="w8a8-mxfp8")
w, sfb = group_wise_quant(w_bf16.view(num_experts * n, k), mode="w8a8-mxfp8")
w = w.view(num_experts, n, k)
sfb = sfb.view(num_experts, n, -1)
# first 128 rows -> expert 0, next 128 -> expert 1
expert_ids = torch.arange(m, device="cuda", dtype=torch.int32) // 128
out = torch.empty(m, n, dtype=torch.bfloat16, device="cuda")
group_contiguous_gemm(x, w, sfa, sfb, expert_ids, out, mode="w8a8-mxfp8")

group_masked_gemm¶

group_masked_gemm( input, weight, sfa, sfb, masked_m, output, bias=None, alpha=None, max_m_per_expert=0, sm_count=-1, mode='w8a8-mxfp8', )¶

Grouped masked GEMM for MoE on Blackwell with per-expert token masking.

Tokens are pre-arranged in a padded [num_experts * max_m_per_expert, K] layout. masked_m specifies how many rows are valid per expert, enabling early-exit for padded rows.

Parameters:

input (torch.Tensor) –
[num_experts * max_m_per_expert, K] CUDA tensor.
- w8a8-mxfp8: dtype float8_e4m3fn.
- w4a4-*: dtype uint8 (packed FP4, logical K is size(1)*2).
weight (torch.Tensor) – [num_experts, N, K] 3D CUDA tensor. Per-expert weights.
sfa (torch.Tensor) –
Scale factor for A, dtype uint8, CUTLASS SfAtom 128x4 swizzled layout. Shape: [num_experts * max_m_per_expert, sf_k], where sf_k = K // 32 (w8a8-mxfp8, w4a4-mxfp4) or K // 16 (w4a4-nvfp4).
- w8a8-mxfp8, w4a4-mxfp4: UE8M0 encoding.
- w4a4-nvfp4: UE4M3 encoding.
sfb (torch.Tensor) – Scale factor for B, dtype uint8, same SfAtom layout as sfa, shape [num_experts, N, sf_k].
masked_m (torch.Tensor) – [num_experts] int32. Actual valid token count per expert.
output (torch.Tensor) – [num_experts * max_m_per_expert, N] bfloat16 output. Written in-place.
bias (torch.Tensor | None) – Optional [num_experts, N] bfloat16 bias. Only supported in w8a8-mxfp8 mode.
alpha (torch.Tensor | None) – Optional float32 per-expert scale with shape [num_experts]. Required for w4a4-nvfp4; not accepted for w4a4-mxfp4.
max_m_per_expert (int) – Maximum tokens per expert (padding bound), must be > 0.
sm_count (int) – Number of SMs to use (-1 for all available).
mode (str) – "w8a8-mxfp8", "w4a4-mxfp4" or "w4a4-nvfp4".

Returns:

None. The result is written in-place to output.

Example

from halops import group_masked_gemm, group_wise_quant

num_experts, max_m_per_expert, n, k = 4, 128, 256, 256
m = num_experts * max_m_per_expert
x_bf16 = torch.randn(m, k, dtype=torch.bfloat16, device="cuda")
w_bf16 = torch.randn(num_experts, n, k, dtype=torch.bfloat16, device="cuda")
x, sfa = group_wise_quant(x_bf16, mode="w8a8-mxfp8")
w, sfb = group_wise_quant(w_bf16.view(num_experts * n, k), mode="w8a8-mxfp8")
w = w.view(num_experts, n, k)
sfb = sfb.view(num_experts, n, -1)
# valid rows of each expert's 128-row slot
masked_m = torch.tensor([128, 64, 0, 32], dtype=torch.int32, device="cuda")
out = torch.empty(m, n, dtype=torch.bfloat16, device="cuda")
group_masked_gemm(
    x, w, sfa, sfb, masked_m, out,
    max_m_per_expert=max_m_per_expert, mode="w8a8-mxfp8",
)

group_wise_quant¶

group_wise_quant( input, mode='w8a8-mxfp8', use_swizzle=True, use_pdl=True, )¶

Group-wise quantization from BF16/FP16 to FP8/FP4 with scale factors.

Quantizes a 2D BF16/FP16 tensor per group along K, producing both the quantized data and the block-scale tensor. Supports FP8 (W8A8-MXFP8) and FP4 (W4A4) quantization. The FP4 recipe is selected explicitly by mode.

The output pairs feed dense_gemm / batch_gemm / group_masked_gemm / group_contiguous_gemm of the same mode.

Parameters:

input (torch.Tensor) – [m, k] CUDA tensor, dtype bfloat16 or float16. Source tensor.
mode (str) – "w8a8-mxfp8", "w4a4-mxfp4" or "w4a4-nvfp4". Selects both the value encoding and the scale recipe.
use_swizzle (bool) – True to output scales in the CUTLASS SfAtom 128x4 interleaved layout consumed by the GEMM kernels; False for a plain row-major scale layout.
use_pdl (bool) – Enable Programmatic Dependent Launch (on by default).

Returns:

A tuple (output_quant, output_sf) where –

output_quant:
- FP8: [m, k] float8_e4m3fn
- FP4: [m, k/2] uint8 (packed nibbles, even index = low nibble)
output_sf: [ceil(m/128)*128, k / group_size] uint8
- w8a8-mxfp8, w4a4-mxfp4: UE8M0 encoding (power-of-2 biased exponent)
- w4a4-nvfp4: UE4M3 encoding (fp8_e4m3fn bit pattern as unsigned byte)

Return type:

Tuple[torch.Tensor, torch.Tensor]

Example

from halops import group_wise_quant

m, k = 256, 512
x_bf16 = torch.randn(m, k, dtype=torch.bfloat16, device="cuda")
# FP8 (W8A8-MXFP8) quantization
x_fp8, sf = group_wise_quant(x_bf16, mode="w8a8-mxfp8")
# FP4 with UE8M0 scales, one per 32 elements (MXFP4)
x_mxfp4, sf = group_wise_quant(x_bf16, mode="w4a4-mxfp4")
# FP4 with UE4M3 scales, one per 16 elements (NVFP4)
x_nvfp4, sf = group_wise_quant(x_bf16, mode="w4a4-nvfp4")

prefill_permute_to_masked_gemm¶

prefill_permute_to_masked_gemm( x, scale, topk_ids, num_local_experts, max_token_num, )¶

Permute tokens into padded expert layout for masked GEMM.

Rearranges tokens from a sorted/topk-selected layout into [E, max_token_num, hidden] padded format suitable for group_masked_gemm, and simultaneously converts scale factors from linear row-major into CUTLASS SfAtom layout.

Parameters:

x (torch.Tensor) – [num_tokens, hidden] CUDA tensor. dtype float8_e4m3fn or bfloat16.
scale (torch.Tensor) – [num_tokens, packed_hidden_scale] CUDA tensor. dtype int32 (broadcast mode): packed_hidden_scale = hidden/128/4, each int32 holds 4 per-group-128 scale bytes. dtype uint8 (direct mode): packed_hidden_scale = hidden/128*4, each 4 consecutive bytes form one uint32 store.
topk_ids (torch.Tensor) – [num_tokens, topk] int64. Expert assignment per token; -1 means unassigned/padding.
num_local_experts (int) – Number of local experts (E).
max_token_num (int) – Maximum tokens per expert. Must be >= 128 and divisible by 128. Must also be large enough to hold the tokens routed to a single expert; the kernel does not clamp, so an expert receiving more than max_token_num tokens writes out of bounds (illegal memory access). Size it with headroom above num_tokens * topk / num_local_experts.

Returns:

A tuple (permute_x, output_scale, permuted_indice_map, token_nums_per_expert) –

permute_x: [E * max_token_num, hidden], same dtype as x.
output_scale: [E, max_token_num, kb_dim] uint8, SfAtom layout (128-byte aligned).
kb_dim = hidden / 32 (both scale modes).
permuted_indice_map: [num_tokens, topk] int32, maps (token, slot) to expert-local offset.
token_nums_per_expert: [E, 1] int32, count of tokens assigned to each expert.

Return type:

Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]

Example

from halops import prefill_permute_to_masked_gemm

num_tokens, hidden, topk = 1024, 2048, 8
num_local_experts = 16
# average tokens per expert is num_tokens * topk / E = 512, so
# max_token_num must leave headroom for routing imbalance
max_token_num = 1024
x = torch.randn(num_tokens, hidden,
                dtype=torch.bfloat16, device="cuda")
# broadcast mode: int32 scale, hidden / 128 / 4 columns
scale = torch.zeros(num_tokens, hidden // 128 // 4,
                    dtype=torch.int32, device="cuda")
topk_ids = torch.randint(0, num_local_experts,
                         (num_tokens, topk),
                         dtype=torch.int64, device="cuda")
permute_x, out_scale, indice_map, token_nums = (
    prefill_permute_to_masked_gemm(
        x, scale, topk_ids, num_local_experts, max_token_num))