every fp8 dim spans at least 2 blocks (grouped kernel)

config.json

@@ -6,13 +6,13 @@
   "first_k_dense_replace": 0,
   "head_dim": 64,
   "hidden_act": "silu",
-  "hidden_size": 64,
+  "hidden_size": 128,
   "initializer_range": 0.02,
-  "intermediate_size": 128,
-  "kv_lora_rank": 32,
+  "intermediate_size": 256,
+  "kv_lora_rank": 64,
   "max_position_embeddings": 1048576,
   "model_type": "mistral4",
-  "moe_intermediate_size": 32,
+  "moe_intermediate_size": 64,
   "n_group": 1,
   "n_routed_experts": 4,
   "n_shared_experts": 1,
@@ -24,7 +24,7 @@
   "output_router_logits": false,
   "pad_token_id": 11,
   "pretraining_tp": 1,
-  "q_lora_rank": 32,
+  "q_lora_rank": 64,
   "qk_head_dim": 64,
   "qk_nope_head_dim": 32,
   "qk_rope_head_dim": 32,

dense_reference.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:eadf4fa154cdf1c8583a58a5d5b8c84ddaec5d635fe3bd7af2a88cf9ede27738
-size 242552
+oid sha256:2ccbf01124aa3f6ea1eeda8f4149c6c3b0d0d56255c2d50ffdf324ddd6639a17
+size 752376

make.py

@@ -16,9 +16,9 @@ from transformers.integrations.nvfp4 import FP4_E2M1_VALUES, SF_VEC_SIZE, NVFP4D
 out = sys.argv[1]
 torch.manual_seed(0)
 config = Mistral4Config(
-    vocab_size=128, hidden_size=64, intermediate_size=128, moe_intermediate_size=32, num_hidden_layers=2,
+    vocab_size=128, hidden_size=128, intermediate_size=256, moe_intermediate_size=64, num_hidden_layers=2,
     num_attention_heads=2, num_key_value_heads=2, n_routed_experts=4, num_experts_per_tok=2, n_shared_experts=1,
-    q_lora_rank=32, kv_lora_rank=32, qk_rope_head_dim=32, qk_nope_head_dim=32, v_head_dim=32, first_k_dense_replace=0,
+    q_lora_rank=64, kv_lora_rank=64, qk_rope_head_dim=32, qk_nope_head_dim=32, v_head_dim=32, first_k_dense_replace=0,
 )  # fmt: skip
 model = Mistral4ForCausalLM(config).to(torch.bfloat16).eval()
 

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:372aa45f180c2fa61ed8fd20177d48749f780ce4e05fbc5b13756613ef3f0e1b
-size 134512
+oid sha256:40890f26eb9448615a409bd6f6b189e8014514bc5dbfce1174558cbda14381d7
+size 391480