Skip to content

Conversation

@kyuyeunk
Copy link
Collaborator

@kyuyeunk kyuyeunk commented Sep 19, 2025

Description

Update tile tuning key for future mixed precision support. Note that this PR does not introduce any new features.

  • add argument x_q_dtype that specifies the precision of activation quantization
  • add extra field to kernel tile tuning key specifying activation quantization dtype
  • make tile tuning key/value data structure flexible by introducing classes TunedKey and TunedValue
  • general formatting & readability improvements

Tests

pytest -v tests/kernels/quantized_matmul_kernel_test.py

https://buildkite.com/tpu-commons/tpu-commons-ci/builds/3156

Checklist

Before submitting this PR, please make sure:

  • I have performed a self-review of my code.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have made or will make corresponding changes to any relevant documentation.

@github-actions
Copy link

Description

Start with a short description of what the PR does and how this is a change from
the past.

The rest of the description includes relevant details and context, examples:

  • why is this change being made,
  • the problem being solved and any relevant context,
  • why this is a good solution,
  • some information about the specific implementation,
  • shortcomings of the solution and possible future improvements.

If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/123456
FIXES: #123456

Tests

Please describe how you tested this change, and include any instructions and/or
commands to reproduce.

Checklist

Before submitting this PR, please make sure:

  • I have performed a self-review of my code.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have made or will make corresponding changes to any relevant documentation.

@kyuyeunk kyuyeunk changed the title [Kernel] Update tile tuning key for future mixed precision support. [Kernel] Update tile tuning key for future mixed precision support Sep 19, 2025
@kyuyeunk kyuyeunk requested a review from hfan September 19, 2025 07:31
@kyuyeunk kyuyeunk merged commit b826f44 into main Sep 19, 2025
2 checks passed
@kyuyeunk kyuyeunk deleted the update_w8a8_kernel branch September 20, 2025 02:16
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants