[Kernel] Honor leading strides in CUTLASS 3.x scaled_mm

[jackLei0901]

Summary

Fix CUTLASS 3.x scaled GEMM addressing for aligned sliced or padded tensor
views by forwarding their actual leading strides instead of reconstructing
tightly packed strides from shape alone.

The operator entry point already accepts row-major A/output and column-major B
with padded leading dimensions. The C3x caller discarded that metadata, which
silently read or wrote incorrect addresses. This fixes A, B, C, and D leading
strides without adding tensor copies; packed inputs retain the same stride
values as before.

Fixes #55534.

Regression coverage

The former combined subset test is split into independent padded-A, padded-B,
and padded-output cases. Keeping the other operands packed in each case makes
the failed stride contract explicit.

Validation

Environment: NVIDIA H800 PCIe (SM90), CUDA 13.0, PyTorch 2.13.0+cu130.

# Five-way baseline/fix matrix (A, B, and output leading strides)
baseline: 8.10% to 8.16% mismatches for each independently padded tensor
patched:  0 mismatches in all five cases; max abs error ~= 3.5e-7

# Modified regression test
PYTHONPATH=$PWD /root/autodl-tmp/vllm-fp8-venv/bin/python -m pytest \
  tests/kernels/quantization/test_cutlass_scaled_mm.py::test_cutlass_strided_subsets \
  -q --disable-warnings
3 passed in 1.91s

# Full CUTLASS scaled-mm test file
PYTHONPATH=$PWD /root/autodl-tmp/vllm-fp8-venv/bin/python -m pytest \
  tests/kernels/quantization/test_cutlass_scaled_mm.py -q --disable-warnings
669 passed, 54 skipped in 1090.56s

# Targeted formatting/lint hooks
uvx pre-commit run clang-format --files \
  csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh
uvx pre-commit run ruff-check --files \
  tests/kernels/quantization/test_cutlass_scaled_mm.py
uvx pre-commit run ruff-format --files \
  tests/kernels/quantization/test_cutlass_scaled_mm.py
All passed

Model evaluations are not applicable: this is a low-level kernel addressing
fix, and the focused numerical comparison covers the affected output contract.

Duplicate-work check

I searched open PRs by issue number and by CUTLASS scaled_mm leading stride
and cutlass_gemm_caller stride. No open PR addresses this bug. The keyword
results #33651 and #41834 concern different kernel/backend changes.

AI assistance disclosure

AI assistance was used to help inspect the CUTLASS/PyTorch stride contract,
prepare the patch and tests, and organize the validation evidence. I reviewed
the changed lines and validation results before submission.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Team

Run ID: db4cd104-e923-4935-b3bf-65e64107bbfa

📥 Commits

Reviewing files that changed from the base of the PR and between 1970f3e and 1084637.

📒 Files selected for processing (2)

• csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh
• tests/kernels/quantization/test_cutlass_scaled_mm.py

Included review availability: Your plan provides up to 10 included reviews per hour; 9 remain after this review.

---

📝 Summary

Summary by CodeRabbit

• Bug Fixes

  • Improved quantized matrix multiplication for padded and non-contiguous tensor views by preserving their runtime strides.
  • Ensures calculations correctly handle strided input and output tensors.

• Tests

  • Added coverage for strided subsets across input, weight, and output tensors.
  • Validates results against existing baseline expectations.

Walkthrough

The CUTLASS 3.x GEMM caller now uses runtime leading strides for input and output tensors. A parameterized regression test covers strided views of A, B, and the output tensor through the raw scaled matrix multiplication operator.

Changes

CUTLASS stride propagation

Layer / File(s)	Summary
Forward runtime leading strides csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh	The caller assigns leading strides from a, b, and out to the corresponding CUTLASS stride objects.
Validate strided subsets tests/kernels/quantization/test_cutlass_scaled_mm.py	A parameterized test covers non-contiguous views for a, b, and out through torch.ops._C.cutlass_scaled_mm.

Estimated code review effort: 2 (Simple) | ~10 minutes

Merge Risk: ⚪ Minimal · up to 10846

CUTLASS 3.x scaled matrix multiplication now correctly honors supported padded tensor views while retaining packed-layout behavior. Coverage includes independently padded inputs and output, with no current merge-blocking risk identified.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 2 functions across 1 files. (1 skipped: 1 …	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the main change: honoring leading strides in the CUTLASS 3.x scaled_mm kernel.
Description check	✅ Passed	The description directly explains the stride-addressing bug, the fix, regression coverage, and validation results.
Linked Issues check	✅ Passed	The changes satisfy issue #55534 by forwarding runtime leading strides for A, B, C, and D, preserving packed-input behavior without copies, and adding independent padded A, B, and output regression ca…
Out of Scope Changes check	✅ Passed	The modified kernel code and regression tests are directly related to the linked stride-addressing bug. No unrelated changes are identified.

Full details: Docstring Coverage

Explanation

Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 2 functions across 1 files. (1 skipped: 1 unsupported.)

• Fix all pre-merge checks with AI

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

• Create stacked PR
• Commit on current branch

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}