Skip to content

Fixes for JIT-able grouped_gemm#1872

Merged
phu0ngng merged 11 commits intoNVIDIA:mainfrom
phu0ngng:jittable_ggemm
Jun 12, 2025
Merged

Fixes for JIT-able grouped_gemm#1872
phu0ngng merged 11 commits intoNVIDIA:mainfrom
phu0ngng:jittable_ggemm

Conversation

@phu0ngng
Copy link
Copy Markdown
Collaborator

@phu0ngng phu0ngng commented Jun 12, 2025 •
edited
Loading

Description

  • Fixes for JIT-able grouped_gemm and grouped_quantize
  • Enforcing workspace to be aligned by 256.

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

Checklist:

  • I have read and followed the contributing guidelines
  • 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

@phu0ngng
Copy link
Copy Markdown
Collaborator Author

phu0ngng commented Jun 12, 2025

/te-ci JAX L0

phu0ngng and others added 6 commits June 12, 2025 10:46
@phu0ngng
fixes for jittable grouped_quantize
7dcf03d 
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@phu0ngng
fixes for jittable grouped_gemm
46705e6 
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@phu0ngng
fix contracting_dim for wgrad gemm
aef9b14 
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@phu0ngng
exclude jitting grouped_gemm fp8
b55174c 
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@pre-commit-ci @phu0ngng
[pre-commit.ci] auto fixes from pre-commit.com hooks
8734831 
for more information, see https://pre-commit.ci
@phu0ngng
cleanup
1252436 
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
phu0ngng added 2 commits June 12, 2025 12:16
@phu0ngng
fixes for jitting grouped_dense
880b625 
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@phu0ngng
Merge branch 'main' into jittable_ggemm
c10065a
@phu0ngng
Copy link
Copy Markdown
Collaborator Author

phu0ngng commented Jun 12, 2025

/te-ci JAX L0

huanghua1994
huanghua1994 previously approved these changes Jun 12, 2025
View reviewed changes
@phu0ngng
exclude all jitted usage in tests
a389771 
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@phu0ngng
Copy link
Copy Markdown
Collaborator Author

phu0ngng commented Jun 12, 2025

/te-ci JAX L0

@phu0ngng phu0ngng merged commit ecaf3e2 into NVIDIA:main Jun 12, 2025
23 checks passed
@phu0ngng phu0ngng deleted the jittable_ggemm branch June 12, 2025 23:01
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@mingxu1067 mingxu1067 Awaiting requested review from mingxu1067

1 more reviewer

@huanghua1994 huanghua1994 huanghua1994 approved these changes

Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@phu0ngng @huanghua1994