Skip to content

[kunlunxin] register new ops and update backend configs - #6009

Open
RobertLuobo wants to merge 4 commits into
flagos-ai:masterfrom
RobertLuobo:dev_configs
Open

[kunlunxin] register new ops and update backend configs#6009
RobertLuobo wants to merge 4 commits into
flagos-ai:masterfrom
RobertLuobo:dev_configs

Conversation

@RobertLuobo

Copy link
Copy Markdown
Contributor
  • Register new ops in FULL_CONFIG (atanh, outer, nextafter.out, unbind_copy.int, bitwise_right_shift.Tensor, rnn_relu.input) and add Quantized dispatch key for dequantize.self
  • Drop cummin/sort from CUSTOMIZED_UNUSED_OPS
  • Add benchmark shapes (grouped_mm, conv_transpose2d, nanmedian, per_token_group_quant_fp8) and baddbmm tune config
  • Add gather_block_quantized BLOCK_SIZE_N heuristic (fixes N=0 arange compile error) and tune index_add/index_select heuristics
  • black formatting fix in pointwise_dynamic.py

PR Category

Type of Change

Description

Issue

Progress

  • Change is properly reviewed (1 reviewer required, 2 recommended).
  • Change is responded to an issue.
  • Change is fully covered by a UT.

Performance

- Register new ops in _FULL_CONFIG (atanh_, outer, nextafter.out,
  unbind_copy.int, bitwise_right_shift.Tensor, rnn_relu.input) and add
  Quantized dispatch key for dequantize.self
- Drop cummin/sort from CUSTOMIZED_UNUSED_OPS
- Add benchmark shapes (grouped_mm, conv_transpose2d, nanmedian,
  per_token_group_quant_fp8) and baddbmm tune config
- Add gather_block_quantized BLOCK_SIZE_N heuristic (fixes N=0 arange
  compile error) and tune index_add/index_select heuristics
- black formatting fix in pointwise_dynamic.py
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Comment thread src/flag_gems/__init__.py
("deg2rad_", deg2rad_),
("dequantize", dequantize),
("dequantize.self", dequantize),
("dequantize.self", dequantize, None, (QUANTIZED_DISPATCH_KEY,)),

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Registering dequantize.self on the Quantized dispatch key makes per-channel quantized tensors reach this implementation, but it unconditionally calls the per-tensor-only q_scale() and q_zero_point() APIs, and the kernel applies a single scale/zero point to every element. Please support per-channel qschemes using q_per_channel_scales(), q_per_channel_zero_points(), and q_per_channel_axis(), or preserve the native path for unsupported qschemes. Please also add torch.quantize_per_channel parity tests for multiple axes.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants