NVIDIA / NVIDIA/cutlass

[QST] How to pack int4 tensor correctly in PyTorch

Open
#2,218 5 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

? - Needs Triage inactive-30d inactive-90d question
Dominant language
C++
Stars
10.5k
Forks
2.1k
Avg merge
3d 11h
Merged PRs (30d)
7

Description

I am trying to use the fp8 * int4 GEMM kernel from 55_hopper_int4_fp8_gemm.cu in PyTorch. I have no problem generate the A matrix which is FP8 [m, k], but not sure how to pack the B matrix with int4 correctly. PyTorch doesn't have the native support of int4, so I have to build and pack it correctly into a uint8 tensor. The B matrix is expected to be [k, n] with torch.int4, but after the packing, the tensor will either be [k // 2, n] or [k, n// 2]. I am not sure if that is correct and expected... So the question is:

Is there any document or example about how should I prepare the int4 matrix correctly in PyTorch to be recognized by the above CUDA kernel example?

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with 55_hopper_int4_fp8_gemm.cu and the PyTorch integration points it uses to determine how the kernel interprets the packed B matrix. Document or demonstrate the expected int4 packing, tensor shape, dtype, and layout so a PyTorch user can prepare an input that the example recognizes.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, pytorch
Domain
machine-learning
Issue type
Documentation
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.