Issues

• Attention's structured Mask and return_residuals in the public API

  openxla/tokamax#1669 · Filter by author AshishKumar4 opened on Oct 19, 2026

• Registering third-party implementations, recording fallbacks, and eligibility by device kind

  openxla/tokamax#1668 · Filter by author AshishKumar4 opened on Oct 19, 2026

• Forward-mode differentiation (jax.jvp) through Ops on Op.__call__'s custom_vjp path

  openxla/tokamax#1667 · Filter by author AshishKumar4 opened on Oct 19, 2026

• Public APIs for serializing BoundArguments, the autotuning data type, and device-kind inference

  openxla/tokamax#1666 · Filter by author AshishKumar4 opened on Oct 19, 2026

• Crash on AMD GPUs: float(compute_capability) fails on ROCm "gfx" strings

  openxla/tokamax#1650 · Filter by author FIM43-Redeye opened on Oct 17, 2026

• linear_softmax_cross_entropy TPU backward ~60% slower in 0.0.16 on JAX <= 0.11.2 (single-buffering fallback), ~6% slower with double buffering

  openxla/tokamax#1547 · Filter by author vlasenkoalexey opened on Oct 9, 2026

• PallasTritonFlashAttentionVjp: wrong causal gradients when block_m1 > block_n1 or block_n2 > block_m2

  openxla/tokamax#1494 · Filter by author AshishKumar4 opened on Oct 2, 2026

• PallasTritonGatedLinearUnit tile heuristic ignores per-block shared-memory capacity (RESOURCE_EXHAUSTED on Blackwell parts with <=96 SMs)

  openxla/tokamax#1361 · Filter by author janbernloehr opened on Sep 15, 2026
  1

• Mosaic-GPU flash attention (sm90/sm100) fails to trace when jax_enable_x64 is enabled

  openxla/tokamax#1349 · Filter by author wittenator opened on Sep 11, 2026

• ragged_dot gives nonzero grad to empty groups

  openxla/tokamax#887 · Filter by author renos opened on Apr 28, 2026
  1

• [TPU] FP8 blockwise quantization: block_size vs MXU utilization tradeoff — how to reconcile in Pallas kernels?

  openxla/tokamax#839 · Filter by author pengchengneo opened on Apr 8, 2026

• Mosaic GPU implementation of mamba2 ssd layer

  openxla/tokamax#750 · Filter by author AlexG1105 opened on Mar 5, 2026
  2

• B200 Mosaic ragged_dot VJP unsupported

  openxla/tokamax#628 · Filter by author changlan opened on Feb 11, 2026
  1

• result mismatch for mosaic gpu attention on sm100

  openxla/tokamax#574 · Filter by author qelk123 opened on Jan 31, 2026
  1

• [Feature Request] Support for Custom DeepSeek-V3 Model Sizes (30B & 155B) in Tokamax GMM

  openxla/tokamax#385 · Filter by author bzantium opened on Nov 24, 2025
  1

• splash attention not compatible with shard_map(check_vma=True)

  openxla/tokamax#382 · Filter by author hb-scaler opened on Nov 21, 2025
  1

• Error tuning ragged_dot backward

  openxla/tokamax#375 · Filter by author hb-scaler opened on Nov 20, 2025
  1

• Trying to access flag --tokamax_xxx before flags were parsed

  openxla/tokamax#374 · Filter by author hb-scaler opened on Nov 20, 2025
  1

• When would the performance benchmark is available?

  openxla/tokamax#298 · Filter by author MoFHeka opened on Nov 1, 2025
  2