biondizzle 05cdde1676 Fix wo_a: scatter each group's data at correct offset in padded buffer ...

The grouped GEMM expects each group's tokens at their own offset range:
- Group 0: rows [0, padded_T)
- Group 1: rows [padded_T, 2*padded_T)
- etc.

Previously we wrote all groups' data contiguously starting at row 0,
so group 1+ would read zeros from the padding area. Now we scatter
each group's quantized activation at the correct offset.

Also:
- Size buffer for total_max_rows = padded_max * n_groups
- Use assemble_scales_2d_side for multi-group scale assembly
- Extract output per-group at correct offsets

2026-05-19 02:45:57 +00:00

..

kernel

refactor: copy CuTeDSL kernel into repo with local imports

2026-05-16 02:57:54 +00:00

__init__.py

refactor: copy CuTeDSL kernel into repo with local imports

2026-05-16 02:57:54 +00:00

bridge.py

Fix torch.compile crash: remove threading.Lock from LUT cache path

2026-05-18 20:54:55 +00:00

custom_ops.py

Replace autograd.Function with torch.library.custom_op for Dynamo compat

2026-05-19 01:54:48 +00:00

inverse_rope.py

Replace DeepGEMM fp8_einsum with CuTeDSL NVFP4 for wo_a (o_proj)

2026-05-19 02:36:30 +00:00

moe_pipeline.py

Add pipeline test with real model weights, add swiglu_limit to reference moe_pipeline

2026-05-17 18:07:44 +00:00

nvfp4_linear.py

Replace autograd.Function with torch.library.custom_op for Dynamo compat

2026-05-19 01:54:48 +00:00

runner.py

Fix gs None error in legacy _ensure_stacked path

2026-05-19 02:17:53 +00:00

shared_expert_pipeline.py

Fix torch.compile: use custom autograd Function instead of @torch.compiler.disable

2026-05-18 21:38:28 +00:00

wo_a_grouped_linear.py

Fix wo_a: scatter each group's data at correct offset in padded buffer

2026-05-19 02:45:57 +00:00