[WIP][feat] grouped fp8: add NT/wgrad persistent path and tile-cumsum group lookup for grouped gemm

[kyle-256]

Description

Two commits on the gpt-oss grouped-fp8 GEMM path (FlyDSL + Triton), rebased onto latest main.

1. grouped fp8 NT + wgrad (FlyDSL). Adds a persistent big-N / short-K path for the
NT forward/dgrad, L2-reuse-aware autotune (NT-store unlocks a wide-gm swizzle band,
plus cs_pipe software-pipelining and non-temporal cstore_aux), and a
distribution-invariant wgrad (flat-by-construction candidate set + grid multiplier for
long-K tails). Motivation: the tensorwise grouped path was L2-reuse-bound, and wgrad
throughput drifted with token-distribution skew.

2. grouped fp8 tensorwise: restore tile-cumsum group lookup (Triton). The tensorwise
persistent kernel had regressed to a per-tile O(G) linear scan (2G group_offs loads per
tile) versus the tile prefix-sum in mpo/0.2.0. This restores a register-resident
tile-cumsum lookup: the host computes cumulative tile counts once (single-program
kernel, same stream, no host sync), and the kernel loads it once and does a vectorized
tl.sum lookup. Byte-exact with the previous scan; empty groups are handled naturally.

Fixes # (n/a)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• FlyDSL: persistent big-N/short-K NT forward/dgrad path in gemm_fp8_grouped_kernel.py
• FlyDSL: L2-reuse-aware autotune — wide-gm swizzle band, cs_pipe, non-temporal cstore_aux
• FlyDSL: distribution-invariant wgrad (flat-by-construction + long-K grid multiplier), helper updates in gemm_helper.py
• Triton: restore register-resident tile-cumsum group lookup in the tensorwise persistent
  kernel, replacing the per-tile O(G) group_offs scan (grouped_gemm_fp8_kernel.py)

Checklist:

• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[Copilot]

Pull request overview

This PR enhances the grouped FP8 GEMM implementation across both Triton and FlyDSL paths, focusing on improving persistent-kernel efficiency and making wgrad/autotune behavior less sensitive to token distribution skew.

Changes:

• Triton tensorwise persistent kernel: restore a tile-cumsum (prefix-sum) group lookup to eliminate per-tile O(G) scans.
• FlyDSL grouped NT: add persistent big‑N/short‑K path improvements (L2-reuse-aware autotune, optional cshuffle soft-pipe + non-temporal C store aux).
• FlyDSL wgrad: adjust candidate selection + autotune scoring to be more distribution-invariant (balanced + canonical skew synthetic loads) and improve long‑K tail utilization.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

File	Description
primus_turbo/triton/grouped_gemm/grouped_gemm_fp8_kernel.py	Adds device-side tile-cumsum precompute and register-resident lookup to avoid per-tile group offset scans.
primus_turbo/flydsl/utils/gemm_helper.py	Extends CShuffle epilogue store to support optional depth-2 soft-pipe and non-temporal store modifier.
primus_turbo/flydsl/grouped_gemm/gemm_fp8_grouped_kernel.py	Updates grouped NT + wgrad kernels and autotuning logic to improve persistence, L2 reuse, and distribution invariance.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 5 out of 6 changed files in this pull request and generated 1 comment.

[Copilot]

Pull request overview

Copilot reviewed 10 out of 11 changed files in this pull request and generated no new comments.

Suppressed comments (2)

primus_turbo/triton/grouped_gemm/grouped_gemm_fp8_kernel.py:441

• tile_start is computed via a full-vector compare + tl.where + tl.sum, which adds an extra O(MAX_G_NEXT_POW2) reduction per tile. Since group_idx is already a scalar, this can be loaded directly from tile_cumsum_ptr (same as total_tiles above) to reduce instruction count and register pressure in the hot loop.

        le_mask = (tile_cumsum_arr <= global_tile_id) & (g_idx_arr >= 1) & (g_idx_arr <= G)
        group_idx = tl.sum(le_mask.to(tl.int32))
        # tile_start = tile_cumsum[group_idx]
        tile_start = tl.sum(tl.where(g_idx_arr == group_idx, tile_cumsum_arr, 0))

primus_turbo/pytorch/ops/grouped_gemm_fp8.py:468

• The new K-pad fast path is guarded on K % 128 != 0 and K >= 129, but the existing grouped GEMM FP8 test matrix only exercises 128-aligned K values. That means this new forward/backward path (including the gradient slicing back to the real K) is currently untested.

        # K-pad fast path: a non-128-aligned K makes the fp8 row stride split each vector load across a 128B line, so cast-and-pad both operands to Kp=ceil128(K) (pad cols contribute 0*0=0, GEMM unchanged). Raw NT inputs only.
        _kpad = (
            trans_b
            and not isinstance(a, QuantizedTensor)
            and not isinstance(b, QuantizedTensor)
            and a.dim() == 2
            and b.dim() == 3
            and (a.shape[-1] % 128 != 0)
            and a.shape[-1] >= 129
        )
        if _kpad:
            from primus_turbo.flydsl.grouped_gemm.gemm_fp8_grouped_kernel import (
                grouped_gemm_fp8_tensorwise_flydsl_kernel,
            )
            from primus_turbo.pytorch.kernels.quantization.quantization_impl import (
                quantize_fp8_tensorwise_pad_impl,
            )

            go = group_offs if group_offs is not None else group_offs_from_lens(group_lens)
            fp8_dtype = _get_fp8_dtype(config.format, True)
            a_q, a_sc = quantize_fp8_tensorwise_pad_impl(a, fp8_dtype)  # [M, Kp]
            b_q, b_sc = quantize_fp8_tensorwise_pad_impl(b, fp8_dtype)  # [G, N, Kp]
            out = grouped_gemm_fp8_tensorwise_flydsl_kernel(
                a_q,
                b_q,
                a_sc,
                b_sc,
                go,
                trans_b=True,
                out_dtype=out_dtype,
                num_cu=(num_cu if num_cu is not None else -1),
            )
            ctx.save_for_backward(a_q, b_q, a_sc, b_sc, group_lens, go)
            ctx.trans_a = False
            ctx.trans_b = trans_b
            ctx.use_nt_layout_gemm_in_bwd = False
            ctx.config = config
            ctx.out_dtype = out_dtype
            ctx.num_cu = num_cu
            ctx.k_pad_real = a.shape[-1]
            return out

[Copilot]

Pull request overview

Copilot reviewed 10 out of 11 changed files in this pull request and generated 1 comment.

[Copilot]

Pull request overview

Copilot reviewed 10 out of 11 changed files in this pull request and generated no new comments.

Suppressed comments (3)

csrc/kernels/quantization/quantization_tensorwise.cu:156

• quantize_tensorwise_pad_impl assumes the padded row length Kp is divisible by the chosen vector UNROLL (it computes cols_per_row = Kp / UNROLL and only processes cols_per_row * UNROLL columns). Right now pack_size is chosen based on K only, so if a caller passes a non power-of-two padding_align_size (making Kp not divisible by pack_size/UNROLL), the kernel can leave the last Kp % UNROLL columns unwritten.

    // Kp is a 128-multiple so Kp % pack == 0 for any pack in {8,4,2,1}. Require
    // K % pack == 0 too so the per-row input base (row*K) keeps vector alignment.
    int32_t pack_size = std::min(get_pack_size<FType>(x), get_pack_size<QType>(y));
    while (pack_size > 1 && (K % pack_size != 0)) {
        pack_size /= 2;
    }

csrc/pytorch/quantization/quantization.cpp:56

• padding_align_size is part of the public op API and currently only checked for >= 1, but the pad kernel relies on alignment properties (see quantize_tensorwise_pad_impl). Without constraining this (e.g., power-of-two), callers can select values that produce a Kp not compatible with vector packing, leading to incorrect output.

    PRIMUS_TURBO_CHECK(input.is_contiguous(), "input must be contiguous");
    PRIMUS_TURBO_CHECK(input.dim() >= 1, "input must have at least 1 dim");
    PRIMUS_TURBO_CHECK(padding_align_size >= 1, "padding_align_size must be >= 1");
    auto stream = at::cuda::getCurrentCUDAStream();

    const int64_t K    = input.size(-1);
    const int64_t Kp   = cdiv(K, padding_align_size) * padding_align_size;
    const int64_t rows = input.numel() / std::max<int64_t>(K, 1);

csrc/pytorch/bindings_pytorch.cpp:32

• Changing the default padding_align_size to 128 changes the output shape for any direct caller of torch.ops.primus_turbo_cpp_extension.quantize_fp8_tensorwise(...) that relies on the previous shape-preserving behavior. Consider keeping the default shape-preserving (padding_align_size=1) and letting higher-level wrappers opt into padding explicitly (as this PR already does).

    m.def("quantize_fp8_tensorwise(Tensor input, ScalarType dest_dtype, Tensor? scale_opt=None, "
          "int padding_align_size=128) -> Tensor[]");

[Copilot]

Pull request overview

Copilot reviewed 10 out of 11 changed files in this pull request and generated no new comments.

Suppressed comments (3)

primus_turbo/pytorch/ops/grouped_gemm_fp8.py:432

• This K-pad fast path directly calls the FlyDSL kernel without checking FlyDSL backend support (e.g., grouped fp8 FlyDSL is gated to gfx950 / compute capability >= (9,5) in GroupedGEMMFP8FlyDSLBackend.is_supported). On other GPUs, this branch can run and fail at compile/runtime.

        # K-pad fast path: a non-128-aligned K makes the fp8 row stride split each vector load across a 128B line, so cast-and-pad both operands to Kp=ceil128(K) (pad cols contribute 0*0=0, GEMM unchanged). Raw NT inputs only.
        _kpad = (
            trans_b
            and not isinstance(a, QuantizedTensor)
            and not isinstance(b, QuantizedTensor)

csrc/pytorch/bindings_pytorch.cpp:32

• The new padding_align_size argument defaults to 128, which changes quantize_fp8_tensorwise behavior for any existing callers that don’t pass this arg (output becomes K-padded and shape changes). To preserve backward-compatible, shape-preserving semantics, the default should be 1 and padding callers should explicitly pass 128.

    m.def("quantize_fp8_tensorwise(Tensor input, ScalarType dest_dtype, Tensor? scale_opt=None, "
          "int padding_align_size=128) -> Tensor[]");

primus_turbo/pytorch/ops/grouped_gemm_fp8.py:428

• The new K-padding branch introduces a distinct quantize+pad path and post-backward slicing, but existing FP8 grouped GEMM tests appear to use only K values divisible by 128, so this behavior likely isn’t exercised. Adding at least one test case with trans_b=True and a non-128-aligned K (e.g., K=129/130) would cover both forward correctness and backward slicing.

        # K-pad fast path: a non-128-aligned K makes the fp8 row stride split each vector load across a 128B line, so cast-and-pad both operands to Kp=ceil128(K) (pad cols contribute 0*0=0, GEMM unchanged). Raw NT inputs only.

[Copilot]

Pull request overview

Copilot reviewed 10 out of 11 changed files in this pull request and generated no new comments.

Suppressed comments (4)

primus_turbo/pytorch/ops/grouped_gemm_fp8.py:430

• The new K-pad fast path calls the FlyDSL grouped GEMM kernel unconditionally when _kpad is true. FlyDSL tensorwise grouped GEMM is gfx950-only (see GroupedGEMMFP8FlyDSLBackend.can_handle), so this will raise or miscompile on other devices instead of falling back to the existing quantize+Triton path.

        if _kpad:
            from primus_turbo.flydsl.grouped_gemm.gemm_fp8_grouped_kernel import (
                grouped_gemm_fp8_tensorwise_flydsl_kernel,
            )
            from primus_turbo.pytorch.kernels.quantization.quantization_impl import (

primus_turbo/pytorch/ops/grouped_gemm_fp8.py:420

• This new K-pad branch introduces a separate execution path, but the existing grouped_gemm_fp8 tests only exercise K values that are multiples of 128. That means the _kpad behavior (pad+quantize+kernel + backward slicing) is currently untested.

        # K-pad fast path: a non-128-aligned K makes the fp8 row stride split each vector load across a 128B line, so cast-and-pad both operands to Kp=ceil128(K) (pad cols contribute 0*0=0, GEMM unchanged). Raw NT inputs only.
        _kpad = (
            trans_b
            and not isinstance(a, QuantizedTensor)
            and not isinstance(b, QuantizedTensor)

csrc/pytorch/quantization/quantization_meta.cpp:18

• The runtime quantize_fp8_tensorwise now hard-requires input.is_contiguous(), but the Meta implementation doesn’t check this. That can cause torch.compile/meta shape propagation to succeed and then fail at runtime for non-contiguous inputs.

    PRIMUS_TURBO_CHECK(input.dim() >= 1, "input must have at least 1 dim");
    PRIMUS_TURBO_CHECK(padding_align_size >= 1, "padding_align_size must be >= 1");
    const int64_t K  = input.size(-1);
    const int64_t Kp = ((K + padding_align_size - 1) / padding_align_size) * padding_align_size;

csrc/pytorch/bindings_pytorch.cpp:32

• Changing the schema default to padding_align_size=128 makes quantize_fp8_tensorwise(input, dtype) no longer shape-preserving (it will pad K to a 128-multiple). That’s a breaking behavior change for any callers using the op directly; internal call sites already pass an explicit align size (1 for legacy, 128 for pad).

    m.def("quantize_fp8_tensorwise(Tensor input, ScalarType dest_dtype, Tensor? scale_opt=None, "
          "int padding_align_size=128) -> Tensor[]");

[Copilot]

Pull request overview

Copilot reviewed 10 out of 11 changed files in this pull request and generated no new comments.

Suppressed comments (3)

primus_turbo/flydsl/utils/gemm_helper.py:874

• Same issue as the non-transposed store path: n_valid becomes None when col_safe=True, but is still passed as mask=... to _buffer_ops.buffer_store. Prefer a constant True mask to indicate an unmasked store.

                for i in range_constexpr(4):
                    n = base_row + n_local + i
                    n_valid = None if self.col_safe else n < self.c_cols
                    val = vec_f32[i]

primus_turbo/flydsl/utils/gemm_helper.py:843

• When col_safe=True, this sets col_valid to None and passes it as mask=... to _buffer_ops.buffer_store. There are no other call sites using mask=None in the codebase, so this is likely to break codegen or produce an invalid store mask. If the intent is an always-on mask (no bounds check), use a constant True mask instead.

This issue also appears on line 871 of the same file.

                col = base_col + tj * 16 + self.lane_id % 16
                col_valid = None if self.col_safe else col < self.c_cols
                vec_f32 = Vec(c_frag[self.c_idx_fn(ti, tj)])
                # Whole-fragment scale: the wave-uniform per-tensor scale packs to v_pk_mul_f32, bit-identical to the per-element form.

csrc/pytorch/bindings_pytorch.cpp:32

• The new padding_align_size argument is exposed with a default of 128. Existing callers that still invoke quantize_fp8_tensorwise(input, dtype, scale_opt) (3 args) will now silently get a padded output shape (K->Kp), which is a breaking API/ABI behavior change compared to the previous shape-preserving op. Consider defaulting padding_align_size to 1 (shape-preserving) and having callers explicitly request padding (as the new Python pad wrapper already does).

    m.def("quantize_fp8_tensorwise(Tensor input, ScalarType dest_dtype, Tensor? scale_opt=None, "
          "int padding_align_size=128) -> Tensor[]");

[Copilot]

Pull request overview

Copilot reviewed 21 out of 22 changed files in this pull request and generated 1 comment.

Suppressed comments (2)

primus_turbo/pytorch/ops/grouped_gemm_fp8.py:447

• The K-pad fast path forces the FlyDSL kernel and bypasses grouped_gemm_fp8_impl dispatch (so backend overrides via GlobalBackendManager are ignored, and num_cu uses a different sentinel convention). Since the tensors are already quantized, you can preserve normal dispatch by calling grouped_gemm_fp8_impl here.

            out = grouped_gemm_fp8_tensorwise_flydsl_kernel(
                a_q,
                b_q,
                a_sc,
                b_sc,
                go,
                trans_b=True,
                out_dtype=out_dtype,
                num_cu=(num_cu if num_cu is not None else -1),
            )

primus_turbo/pytorch/ops/grouped_gemm_fp8.py:425

• This adds a new behavior path for tensorwise grouped GEMM when K % 128 != 0 (quantize+pad on fwd, then slice grads on bwd), but the existing grouped GEMM FP8 tests only cover K values that are multiples of 128. Add at least one test case (e.g. K=1409, trans_b=True) to ensure the K-pad path is exercised and grad shapes/results remain correct.

        # K-pad fast path: a non-128-aligned K makes the fp8 row stride split each vector load across a 128B line, so cast-and-pad both operands to Kp=ceil128(K) (pad cols contribute 0*0=0, GEMM unchanged). Raw NT inputs only.
        _kpad = (
            trans_b
            and not isinstance(a, QuantizedTensor)
            and not isinstance(b, QuantizedTensor)
            and a.dim() == 2
            and b.dim() == 3
            and (a.shape[-1] % 128 != 0)
            and a.shape[-1] >= 129
        )

[Copilot]

Pull request overview

Copilot reviewed 10 out of 11 changed files in this pull request and generated 1 comment.

Suppressed comments (2)

primus_turbo/triton/grouped_gemm/grouped_gemm_fp8_kernel.py:605

• After adding MAX_G_NEXT_POW2 as a tl.constexpr to _compute_tile_cumsum_kernel, the launch needs to pass it as a meta-parameter. Otherwise Triton compilation will fail due to the missing compile-time constant for tl.static_range/tl.arange sizing.

    # Precompute per-group tile cumulative sums on the device (single-program kernel,
    # same CUDA stream → no host sync). The main kernel then loads this tiny array once
    # into registers and skips the per-tile O(G) scan over group_offs.
    num_pid_n_int = (N + blk_n - 1) // blk_n
    tile_cumsum = torch.empty(G + 1, device=a.device, dtype=torch.int32)
    _compute_tile_cumsum_kernel[(1,)](
        group_offs,
        tile_cumsum,
        G,
        num_pid_n_int,
        BLOCK_SIZE_M=blk_m,
    )
    max_g_next_pow2 = _next_pow2(G + 1)

csrc/pytorch/bindings_pytorch.cpp:32

• The Torch op schema sets padding_align_size=128 as the default. This changes the behavior of existing callers that invoke torch.ops.primus_turbo_cpp_extension.quantize_fp8_tensorwise(input, dtype, scale) (previously shape-preserving) to now return a K-padded output by default.

If backward compatibility is desired, consider defaulting padding_align_size to 1 (shape-preserving), and have callers that want padding explicitly pass 128.

    m.def("quantize_fp8_tensorwise(Tensor input, ScalarType dest_dtype, Tensor? scale_opt=None, "
          "int padding_align_size=128) -> Tensor[]");

[Copilot]

Pull request overview

Copilot reviewed 10 out of 11 changed files in this pull request and generated no new comments.

Suppressed comments (2)

primus_turbo/triton/grouped_gemm/grouped_gemm_fp8_kernel.py:347

• _compute_tile_cumsum_kernel uses for g in range(G) but G is a runtime Triton scalar here, which Triton cannot use as a Python range bound (this will fail to JIT/compile). Make G a tl.constexpr (or rewrite using tl.static_range) so the loop can be unrolled.

def _compute_tile_cumsum_kernel(
    group_offs_ptr,  # [G+1] int64
    tile_cumsum_ptr,  # [G+1] int32 (output)
    G,  # runtime number of groups
    num_pid_n,  # runtime int = cdiv(N, BLOCK_SIZE_N)
    BLOCK_SIZE_M: tl.constexpr,
):
    """Single-program kernel; precomputes per-group cumulative tile counts."""
    cumsum: tl.int32 = 0
    tl.store(tile_cumsum_ptr, 0)
    for g in range(G):
        m_g = (tl.load(group_offs_ptr + g + 1) - tl.load(group_offs_ptr + g)).to(tl.int32)
        tiles_g = tl.cdiv(m_g, BLOCK_SIZE_M) * num_pid_n
        cumsum += tiles_g
        tl.store(tile_cumsum_ptr + g + 1, cumsum)

primus_turbo/pytorch/ops/grouped_gemm_fp8.py:455

• The new K-pad fast path is only taken when K % 128 != 0 (and K >= 129), but the existing grouped GEMM FP8 tests currently use only 128-aligned K values, so this branch is likely untested. Please add/extend a unit test that exercises trans_b=True with a non-128-aligned K (e.g. K=129 or K=256+64) and validates forward+backward shapes/accuracy.

        # K-pad fast path: a non-128-aligned K makes the fp8 row stride split each vector load across a 128B line, so cast-and-pad both operands to Kp=ceil128(K) (pad cols contribute 0*0=0, GEMM unchanged). Raw NT inputs only.
        _kpad = (
            trans_b
            and not isinstance(a, QuantizedTensor)
            and not isinstance(b, QuantizedTensor)
            and a.dim() == 2
            and b.dim() == 3
            and (a.shape[-1] % 128 != 0)
            and a.shape[-1] >= 129
        )
        if _kpad:
            from primus_turbo.flydsl.grouped_gemm.grouped_gemm_fp8_kernel import (
                grouped_gemm_fp8_tensorwise_flydsl_kernel,
            )
            from primus_turbo.pytorch.kernels.quantization.quantization_impl import (
                quantize_fp8_tensorwise_pad_impl,
            )

            go = group_offs if group_offs is not None else group_offs_from_lens(group_lens)
            fp8_dtype = _get_fp8_dtype(config.format, True)
            a_q, a_sc = quantize_fp8_tensorwise_pad_impl(a, fp8_dtype)  # [M, Kp]
            b_q, b_sc = quantize_fp8_tensorwise_pad_impl(b, fp8_dtype)  # [G, N, Kp]
            out = grouped_gemm_fp8_tensorwise_flydsl_kernel(
                a_q,
                b_q,
                a_sc,
                b_sc,
                go,
                trans_b=True,
                out_dtype=out_dtype,
                num_cu=(num_cu if num_cu is not None else -1),
            )
            ctx.save_for_backward(a_q, b_q, a_sc, b_sc, group_lens, go)
            ctx.trans_a = False
            ctx.trans_b = trans_b
            ctx.config = config
            ctx.out_dtype = out_dtype
            ctx.num_cu = num_cu
            ctx.k_pad_real = a.shape[-1]
            return out

[Copilot]

Pull request overview

Copilot reviewed 10 out of 11 changed files in this pull request and generated no new comments.

Suppressed comments (3)

primus_turbo/pytorch/ops/grouped_gemm_fp8.py:420

• This K-pad fast path forces the FlyDSL kernel, but it isn’t gated by hardware capability. grouped_gemm_fp8_tensorwise_flydsl_kernel is gfx950/CDNA4-only (see GroupedGEMMFP8FlyDSLBackend.can_handle), so on other GPUs this branch can raise at import-time or fail at runtime even though the regular Triton path would work.

Gate _kpad on is_gfx950() (or add a safe fallback to the Triton implementation when FlyDSL can’t run).

        # K-pad fast path: a non-128-aligned K makes the fp8 row stride split each vector load across a 128B line, so cast-and-pad both operands to Kp=ceil128(K) (pad cols contribute 0*0=0, GEMM unchanged). Raw NT inputs only.
        _kpad = (
            trans_b
            and not isinstance(a, QuantizedTensor)
            and not isinstance(b, QuantizedTensor)

csrc/pytorch/quantization/quantization.cpp:51

• padding_align_size is now a public argument, but the pad kernel assumes the padded K (Kp) is divisible by the chosen vectorization (UNROLL). With the current checks, callers can pass values like 3/65 and silently leave the last Kp % UNROLL columns unwritten because cols_per_row = Kp / UNROLL truncates in quantize_tensorwise_pad_kernel.

Add a stricter validation here (e.g., require padding_align_size == 1 or padding_align_size % 8 == 0) so Kp is always a multiple of any UNROLL in {1,2,4,8}.

    PRIMUS_TURBO_CHECK(input.is_contiguous(), "input must be contiguous");
    PRIMUS_TURBO_CHECK(input.dim() >= 1, "input must have at least 1 dim");
    PRIMUS_TURBO_CHECK(padding_align_size >= 1, "padding_align_size must be >= 1");

csrc/pytorch/quantization/quantization_meta.cpp:16

• The meta implementation should mirror the runtime constraints on padding_align_size. Without validating the alignment here too, shape inference will succeed for values that the real kernel cannot safely handle (e.g., when Kp is not divisible by the chosen vectorization width).

    PRIMUS_TURBO_CHECK(input.dim() >= 1, "input must have at least 1 dim");
    PRIMUS_TURBO_CHECK(padding_align_size >= 1, "padding_align_size must be >= 1");